fix(syntax tree): Fixed typing issue with GeometryTree causing exception

Fixed geometry Tree with actual type enforcement
idaholab · Mar 16, 2024 · 7a24ac0 · 7a24ac0
2 parents ab21045 + 238176b
commit 7a24ac0
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 17 deletions.
diff --git a/montepy/__init__.py b/montepy/__init__.py
@@ -36,6 +36,7 @@
     except ImportError:
         __version__ = "Undefined"
 
+
 # enable deprecated warnings for users
 if not sys.warnoptions:
     import os, warnings

diff --git a/montepy/input_parser/cell_parser.py b/montepy/input_parser/cell_parser.py
@@ -75,7 +75,7 @@ def union(self, p):
     def geometry_expr(self, p):
         left = p.geometry_expr
         right = p.geometry_term
-        nodes = {"left": left.nodes, "operator": p.union, "right": right.nodes}
+        nodes = {"left": left, "operator": p.union, "right": right}
         return syntax_node.GeometryTree("union", nodes, ":", left, right)
 
     @_("geometry_term")
@@ -112,7 +112,7 @@ def geometry_term(self, p):
         for node in node_iter:
             new_tree = syntax_node.GeometryTree(
                 "intersection",
-                {"left": left, "operator": None, "right": node},
+                {"left": left, "operator": syntax_node.PaddingNode(), "right": node},
                 "*",
                 left,
                 node,
@@ -142,7 +142,10 @@ def geometry_factor(self, p):
 
     @_("COMPLEMENT geometry_factory")
     def geometry_factor(self, p):
-        nodes = {"operator": p.COMPLEMENT, "left": p.geometry_factory}
+        nodes = {
+            "operator": syntax_node.PaddingNode(p.COMPLEMENT),
+            "left": p.geometry_factory,
+        }
         return syntax_node.GeometryTree(
             "complement",
             nodes,

diff --git a/montepy/input_parser/syntax_node.py b/montepy/input_parser/syntax_node.py
@@ -222,6 +222,7 @@ class GeometryTree(SyntaxNodeBase):
 
     def __init__(self, name, tokens, op, left, right=None):
         super().__init__(name)
+        assert all(list(map(lambda v: isinstance(v, SyntaxNodeBase), tokens.values())))
         self._nodes = tokens
         self._operator = Operator(op)
         self._left_side = left
@@ -241,19 +242,8 @@ def format(self):
 
     @property
     def comments(self):
-        for key, node in self.nodes.items():
-            if isinstance(node, SyntaxNodeBase):
-                yield from node.comments
-            elif isinstance(node, dict):
-                yield from GeometryTree._recurse_comments(node)
-
-    @staticmethod
-    def _recurse_comments(tree):
-        for node in tree.values():
-            if isinstance(node, SyntaxNodeBase):
-                yield from node.comments
-            elif isinstance(node, dict):
-                yield from GeometryTree._recurse_comments(node)
+        for node in self.nodes.values():
+            yield from node.comments
 
     @property
     def left(self):

diff --git a/tests/test_edge_cases.py b/tests/test_edge_cases.py
@@ -176,3 +176,4 @@ def test_geometry_comments(self):
         cell = montepy.Cell(input)
         # this step caused an error for #163
         cell.comments
+        cell._tree.format()
diff --git a/tests/test_geometry.py b/tests/test_geometry.py
@@ -10,7 +10,7 @@
 class TestHalfSpaceUnit(TestCase):
     def test_init(self):
         surface = montepy.surfaces.CylinderOnAxis()
-        node = montepy.input_parser.syntax_node.GeometryTree("hi", [], "*", " ", " ")
+        node = montepy.input_parser.syntax_node.GeometryTree("hi", {}, "*", " ", " ")
         half_space = HalfSpace(+surface, Operator.UNION, -surface, node)
         self.assertIs(half_space.operator, Operator.UNION)
         self.assertEqual(half_space.left, +surface)