bug[eve]: Preserve annex None value in visitors (#1508)

GridTools · Mar 21, 2024 · 73f3ea0 · 73f3ea0
1 parent e344afd
commit 73f3ea0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/gt4py/eve/visitors.py b/src/gt4py/eve/visitors.py
@@ -168,7 +168,7 @@ def generic_visit(self, node: concepts.RootNode, **kwargs: Any) -> Any:
                 # access to `new_node.annex` implicitly creates the `__node_annex__` attribute in the property getter
                 new_annex_dict = new_node.annex.__dict__
                 for key in self.PRESERVED_ANNEX_ATTRS:
-                    if value := getattr(old_annex, key, None):
+                    if (value := getattr(old_annex, key, NOTHING)) is not NOTHING:
                         assert key not in new_annex_dict
                         new_annex_dict[key] = value
 

diff --git a/tests/eve_tests/unit_tests/test_visitors.py b/tests/eve_tests/unit_tests/test_visitors.py
@@ -19,12 +19,14 @@
 
 def test_annex_preservation(compound_node: eve.Node):
     compound_node.annex.foo = 1
-    compound_node.annex.bar = 2
+    compound_node.annex.bar = None  # None is easily forgotten so test seperately
+    compound_node.annex.baz = 2
 
     class SampleTranslator(eve.NodeTranslator):
-        PRESERVED_ANNEX_ATTRS = ("foo",)
+        PRESERVED_ANNEX_ATTRS = ("foo", "bar")
 
     translated_node = SampleTranslator().visit(compound_node)
 
     assert translated_node.annex.foo == 1
-    assert not hasattr(translated_node.annex, "bar")
+    assert translated_node.annex.bar is None
+    assert not hasattr(translated_node.annex, "baz")