193 lines
5.9 KiB
Python
193 lines
5.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
AST-based transformer to extract App._render_xxx methods into module-level functions.
|
|
|
|
Transformation:
|
|
class App:
|
|
def _render_xxx(self, ...) -> ReturnType:
|
|
self.foo = bar
|
|
self.baz.qux()
|
|
|
|
Becomes:
|
|
|
|
def render_xxx(app: App, ...) -> ReturnType:
|
|
app.foo = bar
|
|
app.baz.qux()
|
|
|
|
class App:
|
|
def _render_xxx(self, ...) -> ReturnType:
|
|
render_xxx(self, ...)
|
|
"""
|
|
|
|
import ast
|
|
import sys
|
|
import argparse
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
|
|
class RenderMethodTransformer(ast.NodeTransformer):
|
|
def __init__(self, class_name: str = "App", method_prefix: str = "_render"):
|
|
self.class_name = class_name
|
|
self.method_prefix = method_prefix
|
|
self.transformed_methods: list[str] = []
|
|
self.in_target_class = False
|
|
self.current_method: Optional[str] = None
|
|
|
|
def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
|
|
if node.name == self.class_name:
|
|
self.in_target_class = True
|
|
self.current_method = None
|
|
node = self.generic_visit(node)
|
|
self.in_target_class = False
|
|
return node
|
|
return node
|
|
|
|
def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
|
|
if not self.in_target_class:
|
|
return node
|
|
|
|
if not node.name.startswith(self.method_prefix):
|
|
return node
|
|
|
|
if node.name == "__init__" or node.name.startswith("__"):
|
|
return node
|
|
|
|
self.current_method = node.name
|
|
new_name = node.name[1:] # Remove underscore prefix: _render_xxx -> render_xxx
|
|
|
|
self.transformed_methods.append(node.name)
|
|
|
|
new_func = ast.FunctionDef(
|
|
name=new_name,
|
|
args=ast.arguments(
|
|
posonlyargs=[],
|
|
args=[ast.arg(arg="app", annotation=ast.Name(id=self.class_name, ctx=ast.Load()))] + node.args.args,
|
|
kwonlyargs=node.args.kwonlyargs,
|
|
kw_defaults=node.args.kw_defaults,
|
|
defaults=node.args.defaults,
|
|
),
|
|
body=[self._transform_body(body_item) for body_item in node.body],
|
|
decorator_list=node.decorator_list,
|
|
returns=node.returns,
|
|
type_comment=node.type_comment,
|
|
)
|
|
|
|
delegation_method = ast.FunctionDef(
|
|
name=node.name,
|
|
args=node.args,
|
|
body=[ast.Expr(value=ast.Call(
|
|
func=ast.Name(id=new_name, ctx=ast.Load()),
|
|
args=[ast.Name(id="self", ctx=ast.Load())] + [self._arg_to_expr(arg) for arg in node.args.args],
|
|
keywords=[ast.keyword(arg=arg.arg, value=ast.Name(id=arg.arg, ctx=ast.Load())) for arg in node.args.kwonlyargs]
|
|
))],
|
|
decorator_list=node.decorator_list,
|
|
returns=node.returns,
|
|
type_comment=node.type_comment,
|
|
)
|
|
|
|
return ast.Module(body=[new_func, delegation_method], type_ignores=[])
|
|
|
|
def _arg_to_expr(self, arg: ast.arg) -> ast.expr:
|
|
return ast.Name(id=arg.arg, ctx=ast.Load())
|
|
|
|
def _transform_body(self, node: ast.AST) -> ast.AST:
|
|
return ast.NodeTransformer.generic_visit(self, node)
|
|
|
|
def visit_Name(self, node: ast.Name) -> ast.Name:
|
|
if isinstance(node.ctx, ast.Load) and node.id == "self":
|
|
return ast.Name(id="app", ctx=node.ctx)
|
|
return node
|
|
|
|
def visit_Attribute(self, node: ast.Attribute) -> ast.Attribute:
|
|
if isinstance(node.value, ast.Name) and node.value.id == "self":
|
|
new_value = ast.Name(id="app", ctx=node.value.ctx)
|
|
return ast.Attribute(value=new_value, attr=node.attr, ctx=node.ctx)
|
|
return self.generic_visit(node)
|
|
|
|
|
|
class AppRenderExtractor(ast.NodeVisitor):
|
|
def __init__(self, class_name: str = "App", method_prefix: str = "_render"):
|
|
self.class_name = class_name
|
|
self.method_prefix = method_prefix
|
|
self.in_target_class = False
|
|
self.methods: dict[str, dict] = {}
|
|
|
|
def visit_ClassDef(self, node: ast.ClassDef) -> None:
|
|
if node.name == self.class_name:
|
|
self.in_target_class = True
|
|
self.generic_visit(node)
|
|
self.in_target_class = False
|
|
|
|
def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
|
|
if not self.in_target_class:
|
|
return
|
|
|
|
if not node.name.startswith(self.method_prefix):
|
|
return
|
|
|
|
if node.name == "__init__" or node.name.startswith("__"):
|
|
return
|
|
|
|
args = [arg.arg for arg in node.args.args]
|
|
self.methods[node.name] = {
|
|
"line": node.lineno,
|
|
"end_line": node.end_lineno,
|
|
"args": args,
|
|
"has_return": any(isinstance(n, ast.Return) for n in ast.walk(node)),
|
|
}
|
|
|
|
|
|
def extract_render_methods(source_path: str, class_name: str = "App", method_prefix: str = "_render") -> dict:
|
|
with open(source_path, "r", encoding="utf-8") as f:
|
|
source = f.read()
|
|
|
|
tree = ast.parse(source)
|
|
extractor = AppRenderExtractor(class_name, method_prefix)
|
|
extractor.visit(tree)
|
|
|
|
return extractor.methods
|
|
|
|
|
|
def transform_file(source_path: str, output_path: Optional[str] = None, class_name: str = "App", method_prefix: str = "_render") -> list[str]:
|
|
with open(source_path, "r", encoding="utf-8") as f:
|
|
source = f.read()
|
|
|
|
tree = ast.parse(source)
|
|
|
|
transformer = RenderMethodTransformer(class_name, method_prefix)
|
|
new_tree = transformer.visit(tree)
|
|
|
|
result = ast.unparse(new_tree)
|
|
|
|
if output_path:
|
|
with open(output_path, "w", encoding="utf-8", newline="") as f:
|
|
f.write(result)
|
|
else:
|
|
print(result)
|
|
|
|
return transformer.transformed_methods
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Transform App._render_xxx methods to module-level functions")
|
|
parser.add_argument("source", help="Source Python file to transform")
|
|
parser.add_argument("-o", "--output", help="Output file (default: stdout)")
|
|
parser.add_argument("-c", "--cls", default="App", help="Class name to transform (default: App)")
|
|
parser.add_argument("-p", "--prefix", default="_render", help="Method prefix to extract (default: _render)")
|
|
parser.add_argument("--extract-only", action="store_true", help="Only extract method info, don't transform")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.extract_only:
|
|
methods = extract_render_methods(args.source, args.cls, args.prefix)
|
|
print(f"Found {len(methods)} methods:")
|
|
for name, info in sorted(methods.items(), key=lambda x: x[1]["line"]):
|
|
print(f" {name}: lines {info['line']}-{info['end_line']}, args={info['args']}")
|
|
else:
|
|
transformed = transform_file(args.source, args.output, args.cls, args.prefix)
|
|
print(f"Transformed {len(transformed)} methods: {transformed}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |