WIP

myvenv/lib/python3.10/site-packages/jinja2/visitor.py

@@ -0,0 +1,92 @@
+"""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