namespace-inject: AST foolery

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