cursed-snakes/namespace-inject.py

import ast
import inspect
from textwrap import dedent
from typing import Callable

##################
# CURSED SECTION #
##################


class namespace_inject:
    """
    This is a decorator that injects one or more functions into
    the local namespace of the decorated function.
    """

    # functions to be injected
    local_fns: tuple[Callable, ...]

    def __init__(self, *local_fn: Callable) -> None:
        self.local_fns = local_fn

    @staticmethod
    def _get_func_def(object: Callable) -> tuple[ast.Module, ast.FunctionDef]:
        """
        Get the AST representation and the contained ast.FunctionDef
        of a function
        """
        object_ast = ast.parse(dedent(inspect.getsource(object)))

        return object_ast, next(
            x for x in ast.walk(object_ast)
            if isinstance(x, ast.FunctionDef) and x.name == object.__name__
        )

    def __call__(self, fn: Callable) -> Callable:
        """
        The actual decorator function
        """
        fn_ast, fn_func_def = self._get_func_def(fn)

        # prepend the local functions to the decorated function's body
        fn_func_def.body = [
            self._get_func_def(lfn)[1]
            for lfn in self.local_fns
        ] + fn_func_def.body

        # remove the @namespace_inject decorator from the decorated function
        fn_func_def.decorator_list = [
            decorator for decorator in fn_func_def.decorator_list
            if not (
                isinstance(decorator, ast.Call)
                and isinstance(decorator.func, ast.Name)
                and decorator.func.id == namespace_inject.__name__
            )
        ]

        # recompile and return the decorated function
        if (fn_sourcefile := inspect.getsourcefile(fn)) is None:
            fn_sourcefile = "<string>"

        exec(
            compile(ast.fix_missing_locations(fn_ast), fn_sourcefile, "exec"),
            fn_globals := {},
        )
        return fn_globals[fn.__name__]

##################
# NORMAL SECTION #
##################


def foo() -> None:
    # foo, but in global namespace it uses an unbound variable

    print(f"{pstr = }")  # type:ignore # noqa: F821


def call_foo_global(pstr: str) -> None:
    """
    Just call foo, so global foo is called.
    This should fail: foo does not know about `pstr`.
    """

    foo()


def call_foo_local(pstr: str) -> None:
    """
    Redefine foo locally, then call - local foo is called.
    In local namespace, `pstr` is defined.
    """
    # call_local_foo got its own foo, so it calls locally

    def foo() -> None:
        # foo in local namespace is fine

        print(f"{pstr = }")

    foo()


@namespace_inject(foo)
def call_foo_injected(pstr: str) -> None:
    """
    Inject the global foo into the local namespace.
    This behaves like `call_foo_local`.
    """

    foo()


def func_info(f: Callable) -> None:
    try:
        print(f"{f.__name__} = {f}, {f.__code__.co_varnames = }")
        f("bar")
        f("baz")
    except Exception:
        print(f"Function {f.__name__} is broken.")


if __name__ == "__main__":
    func_info(call_foo_global)
    func_info(call_foo_local)
    func_info(call_foo_injected)