diff --git a/namespace-inject.py b/namespace-inject.py index 7a0308b..fbad954 100644 --- a/namespace-inject.py +++ b/namespace-inject.py @@ -1,5 +1,20 @@ +import ast +import inspect +from textwrap import dedent from types import FunctionType -from typing import Callable +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: @@ -10,6 +25,23 @@ def func1(pstr: str) -> None: 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"), @@ -24,17 +56,12 @@ def shoehorn(f: Callable) -> Callable: @shoehorn def func2(pstr: str) -> None: - foo() - - -def func_info(f: Callable) -> None: - print(f"{f.__name__ } = {f}, {f.__code__.co_varnames = }") - f("bar") - f("baz") + foo() # type:ignore # noqa: F821 if __name__ == "__main__": func_info(func1) + # ast_dump(func1) func_info(func2) import json