snakemake-resources-log-handler-script.py

import inspect
import json
import os
import sys
from snakemake.resources import parse_resources

# Required for this file to be loaded by snakemake.main() as a --log-handler-script.
def log_handler(msg):
    pass

# This is what we're here for!
def update_resources(defaults):
    # Parse standard k=v resource values (e.g. v must be an int).
    parsed_defaults = parse_resources(defaults)

    if not parsed_defaults:
        return

    # Find snakemake.main() in our call stack.
    stack_frames = [f.frame for f in inspect.stack()]
    snakemake_main = next(filter(lambda f: function_name(f) == "snakemake.main", stack_frames), None)

    if snakemake_main:
        # Find snakemake.main()'s local "resources" variable so we can mutate it
        resources = snakemake_main.f_locals.get("resources", None)
        if resources is not None:
            for k, v in parsed_defaults.items():
                # Don't override resources set by --resources or a profile config
                resources.setdefault(k, v)

def function_name(frame):
    return (module_name(frame) or "<main>") + "." + frame.f_code.co_name

def module_name(frame):
    module = inspect.getmodule(frame.f_code)
    return module.__name__ if module else None

def warn(*msg):
    print(*msg, file = sys.stderr)

try:
    SNAKEMAKE_RESOURCES = os.environ.get("SNAKEMAKE_RESOURCES") or "null"
    try:
        update_resources(json.loads(SNAKEMAKE_RESOURCES))

    except json.JSONDecodeError as e:
        warn(f"ignoring SNAKEMAKE_RESOURCES={SNAKEMAKE_RESOURCES!r} because it is not valid JSON: {e}")

except Exception as e:
    # Oh well, this is only best effort.
    warn(f"failed to apply SNAKEMAKE_RESOURCES={SNAKEMAKE_RESOURCES!r}: {e}")