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):\n    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]