rf: Walk all rules, collect all names

bids-standard · Aug 30, 2024 · 2b16673 · 2b16673
1 parent 0f7bfb1
commit 2b16673
Showing 1 changed file with 17 additions and 1 deletion.
diff --git a/tools/schemacode/bidsschematools/tests/test_expressions.py b/tools/schemacode/bidsschematools/tests/test_expressions.py
@@ -5,7 +5,16 @@
 import pytest
 from pyparsing.exceptions import ParseException
 
-from ..expressions import Array, ASTNode, BinOp, Function, Property, RightOp, expression
+from ..expressions import (
+    Array,
+    ASTNode,
+    BinOp,
+    Element,
+    Function,
+    Property,
+    RightOp,
+    expression,
+)
 from ..types import Namespace
 
 
@@ -148,6 +157,13 @@ def find_names(node: Union[ASTNode, str]):
     elif isinstance(node, Array):
         for element in node.elements:
             yield from find_names(element)
+    elif isinstance(node, Element):
+        yield from find_names(node.name)
+        yield from find_names(node.index)
+    elif isinstance(node, (int, float)):
+        return
+    else:
+        raise TypeError(f"Unexpected node type: {node!r}")
 
 
 @find_names.register