visitor.py 3.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. """API for traversing the AST nodes. Implemented by the compiler and
  2. meta introspection.
  3. """
  4. import typing as t
  5. from .nodes import Node
  6. if t.TYPE_CHECKING:
  7. import typing_extensions as te
  8. class VisitCallable(te.Protocol):
  9. def __call__(self, node: Node, *args: t.Any, **kwargs: t.Any) -> t.Any:
  10. ...
  11. class NodeVisitor:
  12. """Walks the abstract syntax tree and call visitor functions for every
  13. node found. The visitor functions may return values which will be
  14. forwarded by the `visit` method.
  15. Per default the visitor functions for the nodes are ``'visit_'`` +
  16. class name of the node. So a `TryFinally` node visit function would
  17. be `visit_TryFinally`. This behavior can be changed by overriding
  18. the `get_visitor` function. If no visitor function exists for a node
  19. (return value `None`) the `generic_visit` visitor is used instead.
  20. """
  21. def get_visitor(self, node: Node) -> "t.Optional[VisitCallable]":
  22. """Return the visitor function for this node or `None` if no visitor
  23. exists for this node. In that case the generic visit function is
  24. used instead.
  25. """
  26. return getattr(self, f"visit_{type(node).__name__}", None) # type: ignore
  27. def visit(self, node: Node, *args: t.Any, **kwargs: t.Any) -> t.Any:
  28. """Visit a node."""
  29. f = self.get_visitor(node)
  30. if f is not None:
  31. return f(node, *args, **kwargs)
  32. return self.generic_visit(node, *args, **kwargs)
  33. def generic_visit(self, node: Node, *args: t.Any, **kwargs: t.Any) -> t.Any:
  34. """Called if no explicit visitor function exists for a node."""
  35. for node in node.iter_child_nodes():
  36. self.visit(node, *args, **kwargs)
  37. class NodeTransformer(NodeVisitor):
  38. """Walks the abstract syntax tree and allows modifications of nodes.
  39. The `NodeTransformer` will walk the AST and use the return value of the
  40. visitor functions to replace or remove the old node. If the return
  41. value of the visitor function is `None` the node will be removed
  42. from the previous location otherwise it's replaced with the return
  43. value. The return value may be the original node in which case no
  44. replacement takes place.
  45. """
  46. def generic_visit(self, node: Node, *args: t.Any, **kwargs: t.Any) -> Node:
  47. for field, old_value in node.iter_fields():
  48. if isinstance(old_value, list):
  49. new_values = []
  50. for value in old_value:
  51. if isinstance(value, Node):
  52. value = self.visit(value, *args, **kwargs)
  53. if value is None:
  54. continue
  55. elif not isinstance(value, Node):
  56. new_values.extend(value)
  57. continue
  58. new_values.append(value)
  59. old_value[:] = new_values
  60. elif isinstance(old_value, Node):
  61. new_node = self.visit(old_value, *args, **kwargs)
  62. if new_node is None:
  63. delattr(node, field)
  64. else:
  65. setattr(node, field, new_node)
  66. return node
  67. def visit_list(self, node: Node, *args: t.Any, **kwargs: t.Any) -> t.List[Node]:
  68. """As transformers may return lists in some places this method
  69. can be used to enforce a list as return value.
  70. """
  71. rv = self.visit(node, *args, **kwargs)
  72. if not isinstance(rv, list):
  73. return [rv]
  74. return rv