以查询取代 AST 节点
原标题：Inline replace of AST node by query
Given input like:
query = ("class_name", "function_name", "arg_name")
How do I replace what is found with some other provided node?
Parsed example of the before stage:
class Foo(object):
def f(self, g: str = "foo"): pass
Parsed example of the after stage:
class Foo(object):
def f(self, g: int = 5): pass
Given the following call to a hypothetical function:
replace_ast_node(
query=("Foo", "f", "g"),
node=ast.parse("class Foo(object):
def f(self, g: str = foo ): pass"),
# Use `AnnAssign` over `arg`; as `defaults` is higher in the `FunctionDef`
replace_with=AnnAssign(
annotation=Name(ctx=Load(), id="int"),
simple=1,
target=Name(ctx=Store(), id="g"),
value=Constant(kind=None, value=5),
),
)
I ve hacked together a simple solution for finding a node with the query list, which has the added benefit of working for anything ("Foo", "f", "g") could refer to def Foo(): def f(): def g():, as well as a parser/emitter from arg to AnnAssign. But I can t figure this stage out; does ast.NodeTransformer traverse in-order, and sequentially? - So should I be constantly traversing, appending current name, and checking if current location is the full query string? - I feel like there s some clean solution I m missing…
最佳回答
I decided to split it into two problems. The first, give every node awareness of its place in the universe:
def annotate_ancestry(node):
"""
Look to your roots. Find the child; find the parent.
Sets _location attribute to every child node.
:param node: AST node. Will be annotated in-place.
:type node: ```ast.AST```
"""
node._location = [node.name] if hasattr(node, name ) else []
parent_location = []
for _node in ast.walk(node):
name = [_node.name] if hasattr(_node, name ) else []
for child in ast.iter_child_nodes(_node):
if hasattr(child, name ):
child._location = name + [child.name]
parent_location = child._location
elif isinstance(child, ast.arg):
child._location = parent_location + [child.arg]
Then implement one method of the aforementioned ast.NodeTransformer:
class RewriteAtQuery(ast.NodeTransformer):
"""
Replace the node at query with given node
:ivar search: Search query, e.g., [ class_name , method_name , arg_name ]
:ivar replacement_node: Node to replace this search
"""
def __init__(self, search, replacement_node):
"""
:param search: Search query
:type search: ```List[str]```
:param replacement_node: Node to replace this search
:type replacement_node: ```ast.AST```
"""
self.search = search
self.replacement_node = replacement_node
self.replaced = False
def generic_visit(self, node):
"""
Visit every node, replace once, and only if found
:param node: AST node
:type node: ```ast.AST```
:returns: AST node, potentially edited
:rtype: ```ast.AST```
"""
if not self.replaced and hasattr(node, _location )
and node._location == self.search:
node = self.replacement_node
self.replaced = True
return ast.NodeTransformer.generic_visit(self, node)
And you re done =)
Usage/test:
parsed_ast = ast.parse(class_with_method_and_body_types_str)
annotate_ancestry(parsed_ast)
rewrite_at_query = RewriteAtQuery(
search="C.method_name.dataset_name".split("."),
replacement_node=arg(
annotation=Name(ctx=Load(), id="int"),
arg="dataset_name",
type_comment=None,
),
).visit(parsed_ast)
self.assertTrue(rewrite_at_query.replaced, True)
# Additional test to compare AST produced with desired AST [see repo]
问题回答
暂无回答
