feat: splice children

mecha/ast.py

@@ -83,6 +83,7 @@
 from typing import (
     Any,
     ClassVar,
+    Iterable,
     Iterator,
     List,
     Literal,
@@ -197,6 +198,18 @@ def emit_error(self, exc: T) -> T:
 class AstChildren(Tuple[AstNodeType, ...]):
     """Specialized tuple subclass for holding multiple child ast nodes."""
 
+    def __new__(
+        cls,
+        children: Iterable[None | AstNodeType | "AstChildren[AstNodeType]"] = (),
+    ):
+        children = [
+            c
+            for child in children
+            if child is not None
+            for c in (child if isinstance(child, AstChildren) else [child])
+        ]
+        return super().__new__(cls, children)
+
     def __repr__(self) -> str:
         return f"AstChildren({super().__repr__()})"
 

mecha/dispatch.py

@@ -373,7 +373,7 @@ def invoke(self, node: AstNode, *args: Any, **kwargs: Any) -> Any:
             attribute = getattr(node, f.name)
             if isinstance(attribute, AstChildren):
                 result = AstChildren(self.invoke(child, *args, **kwargs) for child in attribute)  # type: ignore
-                if any(child is not original for child, original in zip(result, attribute)):  # type: ignore
+                if len(result) != len(attribute) or any(child is not original for child, original in zip(result, attribute)):  # type: ignore
                     to_replace[f.name] = result
             elif isinstance(attribute, AstNode):
                 result = self.invoke(attribute, *args, **kwargs)
@@ -383,7 +383,6 @@ def invoke(self, node: AstNode, *args: Any, **kwargs: Any) -> Any:
         if to_replace:
             node = replace(node, **to_replace)
 
-        result = node
         exhausted = False
 
         while not exhausted:
@@ -397,7 +396,9 @@ def invoke(self, node: AstNode, *args: Any, **kwargs: Any) -> Any:
                         exhausted = False
                         node = result
                         break
-                    elif result is not None:
+                    elif result is None or isinstance(result, AstChildren):
+                        return result  # type: ignore
+                    else:
                         msg = f"Invalid node of type {type(result)}."
                         if name:
                             msg += f" ({name})"