cursed-snakes/namespace-inject.py

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

#################
# MAGIC SECTION #
#################


def ast_dump(sobject) -> None:
    print(f"Dumping AST of {sobject}")
    print(ast.dump(ast.parse(
        dedent(inspect.getsource(sobject)),
    ), indent=4))


def shoehorn(f: Callable) -> Callable:
    # ast_dump(f)

    def foo() -> None:
        print(f"{pstr = }")  # type:ignore # noqa: F821

    # ast_dump(foo)
    foo_ast = ast.parse(dedent(inspect.getsource(foo)))
    foo_fn = next(
        x
        for x in ast.walk(foo_ast)
        if isinstance(x, ast.FunctionDef) and x.name == "foo"
    )

    class Shoehorn(ast.NodeTransformer):
        def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
            if node.name == f.__name__:
                node.body.insert(0, foo_fn)
                node.decorator_list = [
                    decorator
                    for decorator in node.decorator_list
                    if not (
                        isinstance(decorator, ast.Name)
                        and decorator.id == shoehorn.__name__
                    )
                ]
            return node

    f_ast = ast.parse(dedent(inspect.getsource(f)))
    new_f_ast = ast.fix_missing_locations(Shoehorn().visit(f_ast))
    # print(ast.dump(new_f_ast, indent=4))

    new_f_scope = {}
    exec(compile(new_f_ast, "<string>", "exec"), new_f_scope)

    return new_f_scope[f.__name__]

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


def foo():
    raise NotImplementedError()


def func1(pstr: str) -> None:
    def foo() -> None:
        print(f"{pstr = }")

    foo()


def func2(pstr: str) -> None:
    foo()


@shoehorn
def func3(pstr: str) -> None:
    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(func1)
    func_info(func2)
    func_info(func3)