cursed-snakes/namespace-inject.py

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


def ast_dump(sobject) -> None:
    print(ast.dump(ast.parse(
        source=dedent(inspect.getsource(sobject)),
    ), indent=4))


def func_info(f: Callable) -> None:
    print(f"{f.__name__ } = {f}, {f.__code__.co_varnames = }")
    f("bar")
    f("baz")


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

    foo()


def shoehorn(f: Callable) -> Callable:
    ast_dump(f)
    f_ast = ast.parse(source=dedent(inspect.getsource(f)))

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

    ast_dump(foo)
    # foo_ast = ast.parse(source=dedent(inspect.getsource(foo)))

    class Shoehorn(ast.NodeTransformer):
        def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
            print(ast.dump(node, indent=4))
            return node

    new_f_ast = ast.fix_missing_locations(Shoehorn().visit(f_ast))
    print(ast.dump(new_f_ast, indent=4))

    code = f.__code__.replace(
        co_nlocals=2,
        co_varnames=("pstr", "foo"),
    )

    return FunctionType(
        name=f.__name__,
        code=code,
        globals=func1.__globals__,
    )


@shoehorn
def func2(pstr: str) -> None:
    foo()  # type:ignore # noqa: F821


if __name__ == "__main__":
    func_info(func1)
    # ast_dump(func1)
    func_info(func2)

    import json

    print(json.dumps({
        name: repr(getattr(func2.__code__, name))
        for name in dir(func2.__code__)
    }, indent=2))