memory profiler embedded into click command group

click_memory_profiler.py

DEFAULT_DIR = "/tmp/memory_profile"

class RunMprof:

    def __init__(self):
        self.thread: subprocess.Popen = None

    def __enter__(self):
        import datetime, os, pathlib

        output_dir = pathlib.Path(DEFAULT_DIR)
        output_dir.mkdir(exist_ok=True, parents=True)
        timestamp = datetime.datetime.utcnow().strftime("%Y_%m_%d_%H_%m_%S.%f")
        # current process id
        pid_id = os.getpid()
        print(f"pid_id: {pid_id}")
        # NOTE: mprof argument parsing is FUBAR, attach must go last, otherwise output file gets default, and not the user-request
        # value
        cmd =f"mprof run --output {str(output_dir)}/mprofile_{timestamp}.dat --include-children --nopython --attach {pid_id}" 
        print(cmd)
        self.thread = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            shell=True,
        )

    def __exit__(self, exc_type, exc_val, exc_tb):
        # make sure to close subprocess with mprof
        self.thread.terminate()
        self.thread.wait()


@click.group()
@click.pass_context
def cli_driver(ctx):
    import os

    if os.environ.get("DEBUG"):
        ctx.with_resource(RunMprof())

    ctx.obj = get_config()