1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192 |
- """API for traversing the AST nodes. Implemented by the compiler and
- meta introspection.
- """
- import typing as t
- from .nodes import Node
- if t.TYPE_CHECKING:
- import typing_extensions as te
- class VisitCallable(te.Protocol):
- def __call__(self, node: Node, *args: t.Any, **kwargs: t.Any) -> t.Any:
- ...
- class NodeVisitor:
- """Walks the abstract syntax tree and call visitor functions for every
- node found. The visitor functions may return values which will be
- forwarded by the `visit` method.
- Per default the visitor functions for the nodes are ``'visit_'`` +
- class name of the node. So a `TryFinally` node visit function would
- be `visit_TryFinally`. This behavior can be changed by overriding
- the `get_visitor` function. If no visitor function exists for a node
- (return value `None`) the `generic_visit` visitor is used instead.
- """
- def get_visitor(self, node: Node) -> "t.Optional[VisitCallable]":
- """Return the visitor function for this node or `None` if no visitor
- exists for this node. In that case the generic visit function is
- used instead.
- """
- return getattr(self, f"visit_{type(node).__name__}", None) # type: ignore
- def visit(self, node: Node, *args: t.Any, **kwargs: t.Any) -> t.Any:
- """Visit a node."""
- f = self.get_visitor(node)
- if f is not None:
- return f(node, *args, **kwargs)
- return self.generic_visit(node, *args, **kwargs)
- def generic_visit(self, node: Node, *args: t.Any, **kwargs: t.Any) -> t.Any:
- """Called if no explicit visitor function exists for a node."""
- for node in node.iter_child_nodes():
- self.visit(node, *args, **kwargs)
- class NodeTransformer(NodeVisitor):
- """Walks the abstract syntax tree and allows modifications of nodes.
- The `NodeTransformer` will walk the AST and use the return value of the
- visitor functions to replace or remove the old node. If the return
- value of the visitor function is `None` the node will be removed
- from the previous location otherwise it's replaced with the return
- value. The return value may be the original node in which case no
- replacement takes place.
- """
- def generic_visit(self, node: Node, *args: t.Any, **kwargs: t.Any) -> Node:
- for field, old_value in node.iter_fields():
- if isinstance(old_value, list):
- new_values = []
- for value in old_value:
- if isinstance(value, Node):
- value = self.visit(value, *args, **kwargs)
- if value is None:
- continue
- elif not isinstance(value, Node):
- new_values.extend(value)
- continue
- new_values.append(value)
- old_value[:] = new_values
- elif isinstance(old_value, Node):
- new_node = self.visit(old_value, *args, **kwargs)
- if new_node is None:
- delattr(node, field)
- else:
- setattr(node, field, new_node)
- return node
- def visit_list(self, node: Node, *args: t.Any, **kwargs: t.Any) -> t.List[Node]:
- """As transformers may return lists in some places this method
- can be used to enforce a list as return value.
- """
- rv = self.visit(node, *args, **kwargs)
- if not isinstance(rv, list):
- return [rv]
- return rv
|