Repository: GridTools/gt4py
Pull request #1874 (2025-02-19 01:46:47)

feat[eve]: Preserve annex in custom visitors 
Preserve annex attributes in the `NodeTranslator` for nodes returned from custom `visit_` methods, not only from nodes which are recreated in `generic_visit`. This is often desired and with this PR doesn't need to be done manually in child classes of `NodeTranslator`. The `visit` method can still decide if the node it returns should have a different annex attribute as `_preserve_annex` does not overwrite in case the attribute already exists.


Patch:
diff --git a/.python-version b/.default-python-version
similarity index 100%
rename from .python-version
rename to .default-python-version
diff --git a/.github/actions/get-python-versions/action.yml b/.github/actions/get-python-versions/action.yml
new file mode 100644
index 00000000..82d23d82
--- /dev/null
+++ b/.github/actions/get-python-versions/action.yml
@@ -0,0 +1,17 @@
+name: "Get Python Versions"
+description: "Reads Python versions from .python-versions file and outputs them as a JSON array"
+
+outputs:
+  python-versions:
+    description: "JSON array of Python versions"
+    value: ${{ steps.set-versions.outputs.python-versions }}
+
+runs:
+  using: "composite"
+  steps:
+  - uses: actions/checkout@v4
+  - id: set-versions
+    shell: bash
+    run: |
+      PYTHON_VERSIONS=$(cat .python-versions | jq -R -s -c 'split("\n") | map(select(length > 0))')
+      echo "python-versions=${PYTHON_VERSIONS}" >> $GITHUB_OUTPUT
diff --git a/.github/workflows/_disabled/documentation.yml b/.github/workflows/_disabled/documentation.yml
index 71361549..f0bcbecc 100644
--- a/.github/workflows/_disabled/documentation.yml
+++ b/.github/workflows/_disabled/documentation.yml
@@ -9,18 +9,31 @@ on:
     - functional
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.set-matrix.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: set-matrix
+      run: |
+        PYTHON_VERSIONS=$(cat .python-versions | jq -R -s -c 'split("\n") | map(select(length > 0))')
+        echo "python-versions=${PYTHON_VERSIONS}" >> $GITHUB_OUTPUT
+
   quickstart-guide:
+    needs: get-python-versions
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
 
     steps:
     - uses: actions/checkout@v3
-    - name: Set up Python ${{ matrix.python-version }}
+    - name: Set up Python
       uses: actions/setup-python@v4
       with:
-        python-version: ${{ matrix.python-version }}
+        python-version-file: ".default-python-version"
     - name: Install python dependencies
       run: |
         python -m pip install --upgrade pip setuptools wheel
diff --git a/.github/workflows/_disabled/gt4py-sphinx.yml b/.github/workflows/_disabled/gt4py-sphinx.yml
index 2533b2a4..2b071582 100644
--- a/.github/workflows/_disabled/gt4py-sphinx.yml
+++ b/.github/workflows/_disabled/gt4py-sphinx.yml
@@ -20,7 +20,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v1
       with:
-        python-version: 3.10
+        python-version-file: ".default-python-version"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
diff --git a/.github/workflows/_disabled/template_coverage.yml b/.github/workflows/_disabled/template_coverage.yml
index e592ca32..ec0aa436 100644
--- a/.github/workflows/_disabled/template_coverage.yml
+++ b/.github/workflows/_disabled/template_coverage.yml
@@ -16,14 +16,12 @@ on:
 jobs:
   upload:
     runs-on: ubuntu-latest
-    env:
-      PY_VERSION: "3.10"
     steps:
     - uses: actions/checkout@v3
-    - name: Set up Python ${{ env.PY_VERSION }}
+    - name: Set up Python
       uses: actions/setup-python@v4
       with:
-        python-version: ${{ env.PY_VERSION }}
+        python-version-file: ".default-python-version"
 
     - name: Install python dependencies
       shell: bash
diff --git a/.github/workflows/code-quality.yml b/.github/workflows/code-quality.yml
index 10bb537e..09e1ea0a 100644
--- a/.github/workflows/code-quality.yml
+++ b/.github/workflows/code-quality.yml
@@ -17,7 +17,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version-file: ".python-version"
+        python-version-file: ".default-python-version"
 
     - name: Install uv
       uses: astral-sh/setup-uv@v5
diff --git a/.github/workflows/daily-ci.yml b/.github/workflows/daily-ci.yml
index a2a52ce1..1497853e 100644
--- a/.github/workflows/daily-ci.yml
+++ b/.github/workflows/daily-ci.yml
@@ -13,7 +13,18 @@ on:
   ## END
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.get-versions.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: get-versions
+      uses: ./.github/actions/get-python-versions
+
   daily-ci:
+    needs: get-python-versions
     strategy:
       matrix:
         # dependencies-strategy -> The strategy that `uv lock` should use to select
@@ -22,7 +33,7 @@ jobs:
         dependencies-strategy: ["lowest-direct", "highest"]
         gt4py-module: ["cartesian", "eve", "next", "storage"]
         os: ["ubuntu-latest"] #, "macos-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
       fail-fast: false
 
     runs-on: ${{ matrix.os }}
diff --git a/.github/workflows/deploy-release.yml b/.github/workflows/deploy-release.yml
index 7a7505ca..ab7d0343 100644
--- a/.github/workflows/deploy-release.yml
+++ b/.github/workflows/deploy-release.yml
@@ -15,10 +15,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
-    - name: Set up Python 3.10
+    - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version: "3.10"
+        python-version-file: ".default-python-version"
     - name: Install pypa/build
       run: |
         python -m pip install build --user
diff --git a/.python-versions b/.python-versions
new file mode 100644
index 00000000..6db9dc12
--- /dev/null
+++ b/.python-versions
@@ -0,0 +1,2 @@
+3.10
+3.11
diff --git a/noxfile.py b/noxfile.py
index 3aad5658..5fad105c 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -111,7 +111,6 @@ def test_examples(session: nox.Session) -> None:
     session.run(*"jupytext docs/user/next/QuickstartGuide.md --to .ipynb".split())
     session.run(*"jupytext docs/user/next/advanced/*.md --to .ipynb".split())
 
-    num_processes = os.environ.get("NUM_PROCESSES", "auto")
     for notebook, extra_args in [
         ("docs/user/next/workshop/slides", None),
         ("docs/user/next/workshop/exercises", ["-k", "solutions"]),
@@ -120,7 +119,7 @@ def test_examples(session: nox.Session) -> None:
         ("examples", (None)),
     ]:
         session.run(
-            *f"pytest --nbmake {notebook} -sv -n {num_processes}".split(),
+            *f"pytest --nbmake {notebook} -sv -n 1 --benchmark-disable".split(),
             *(extra_args or []),
         )
 
diff --git a/src/gt4py/eve/visitors.py b/src/gt4py/eve/visitors.py
index 59b4ef08..6e559a26 100644
--- a/src/gt4py/eve/visitors.py
+++ b/src/gt4py/eve/visitors.py
@@ -121,6 +121,20 @@ class NodeVisitor:
         return None
 
 
+def _preserve_annex(
+    node: concepts.Node, new_node: concepts.Node, preserved_annex_attrs: tuple[str, ...]
+) -> None:
+    # access to `new_node.annex` implicitly creates the `__node_annex__` attribute in the property getter
+    new_annex_dict = new_node.annex.__dict__
+    old_annex_dict = node.annex.__dict__
+    for key in (old_annex_dict.keys() & preserved_annex_attrs) - new_annex_dict.keys():
+        # Note: The annex item's value of the new node might not be equal
+        # (in the sense that an equality comparison returns false), to
+        # the old one, but in the context of the pass, they are
+        # equivalent. Therefore, we don't assert equality here.
+        new_annex_dict[key] = old_annex_dict[key]
+
+
 class NodeTranslator(NodeVisitor):
     """Special `NodeVisitor` to translate nodes and trees.
 
@@ -158,13 +172,8 @@ class NodeTranslator(NodeVisitor):
                     if (new_child := self.visit(child, **kwargs)) is not NOTHING
                 }
             )
-            if self.PRESERVED_ANNEX_ATTRS and (old_annex := getattr(node, "__node_annex__", None)):
-                # 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, NOTHING)) is not NOTHING:
-                        assert key not in new_annex_dict
-                        new_annex_dict[key] = value
+            if self.PRESERVED_ANNEX_ATTRS and getattr(node, "__node_annex__", None):
+                _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS)
 
             return new_node
 
@@ -189,3 +198,16 @@ class NodeTranslator(NodeVisitor):
             )
 
         return copy.deepcopy(node, memo=memo)
+
+    def visit(self, node: concepts.RootNode, **kwargs: Any) -> Any:
+        new_node = super().visit(node, **kwargs)
+
+        if (
+            isinstance(node, concepts.Node)
+            and isinstance(new_node, concepts.Node)
+            and self.PRESERVED_ANNEX_ATTRS
+            and getattr(node, "__node_annex__", None)
+        ):
+            _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS)
+
+        return new_node
diff --git a/src/gt4py/next/iterator/transforms/infer_domain.py b/src/gt4py/next/iterator/transforms/infer_domain.py
index 239824ba..ee9b695a 100644
--- a/src/gt4py/next/iterator/transforms/infer_domain.py
+++ b/src/gt4py/next/iterator/transforms/infer_domain.py
@@ -22,7 +22,7 @@ from gt4py.next.iterator.ir_utils import (
     ir_makers as im,
     misc as ir_misc,
 )
-from gt4py.next.iterator.transforms import trace_shifts
+from gt4py.next.iterator.transforms import constant_folding, trace_shifts
 from gt4py.next.utils import flatten_nested_tuple, tree_map
 
 
@@ -452,8 +452,12 @@ def _infer_stmt(
     **kwargs: Unpack[InferenceOptions],
 ):
     if isinstance(stmt, itir.SetAt):
+        # constant fold once otherwise constant folding after domain inference might create (syntactic) differences
+        # between the domain stored in IR and in the annex
+        domain = constant_folding.ConstantFolding.apply(stmt.domain)
+
         transformed_call, _ = infer_expr(
-            stmt.expr, domain_utils.SymbolicDomain.from_expr(stmt.domain), **kwargs
+            stmt.expr, domain_utils.SymbolicDomain.from_expr(domain), **kwargs
         )
 
         return itir.SetAt(


Test patch:
diff --git a/.github/workflows/test-cartesian-fallback.yml b/.github/workflows/test-cartesian-fallback.yml
index a846af2e..1cce8831 100644
--- a/.github/workflows/test-cartesian-fallback.yml
+++ b/.github/workflows/test-cartesian-fallback.yml
@@ -12,12 +12,23 @@ on:
     - "*.rst"
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.get-versions.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: get-versions
+      uses: ./.github/actions/get-python-versions
+
   test-cartesian:
+    needs: get-python-versions
     strategy:
       matrix:
         codegen-factor: [internal, dace]
         os: ["ubuntu-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
     runs-on: ${{ matrix.os }}
     steps:
     - run: 'echo "No build required"'
diff --git a/.github/workflows/test-cartesian.yml b/.github/workflows/test-cartesian.yml
index ea6b7940..5c7e2e90 100644
--- a/.github/workflows/test-cartesian.yml
+++ b/.github/workflows/test-cartesian.yml
@@ -19,12 +19,24 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.get-versions.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: get-versions
+      uses: ./.github/actions/get-python-versions
+
+  # Main job that runs the tests
   test-cartesian:
+    needs: get-python-versions
     strategy:
       matrix:
         codegen-factor: [internal, dace]
         os: ["ubuntu-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
       fail-fast: false
 
     runs-on: ${{ matrix.os }}
diff --git a/.github/workflows/test-eve-fallback.yml b/.github/workflows/test-eve-fallback.yml
index f3dbb58a..94e870dc 100644
--- a/.github/workflows/test-eve-fallback.yml
+++ b/.github/workflows/test-eve-fallback.yml
@@ -16,11 +16,22 @@ on:
     - "*.yml"
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.get-versions.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: get-versions
+      uses: ./.github/actions/get-python-versions
+
   test-eve:
+    needs: get-python-versions
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
 
     runs-on: ${{ matrix.os }}
     steps:
diff --git a/.github/workflows/test-eve.yml b/.github/workflows/test-eve.yml
index aad3971a..76bcc972 100644
--- a/.github/workflows/test-eve.yml
+++ b/.github/workflows/test-eve.yml
@@ -16,11 +16,22 @@ on:
     - "*.yml"
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.get-versions.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: get-versions
+      uses: ./.github/actions/get-python-versions
+
   test-eve:
+    needs: get-python-versions
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
       fail-fast: false
 
     runs-on: ${{ matrix.os }}
diff --git a/.github/workflows/test-examples.yml b/.github/workflows/test-examples.yml
index 836af45d..93c703c2 100644
--- a/.github/workflows/test-examples.yml
+++ b/.github/workflows/test-examples.yml
@@ -9,11 +9,22 @@ on:
     - main
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.get-versions.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: get-versions
+      uses: ./.github/actions/get-python-versions
+
   test-notebooks:
+    needs: get-python-versions
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
       fail-fast: false
 
     runs-on: ${{ matrix.os }}
diff --git a/.github/workflows/test-next-fallback.yml b/.github/workflows/test-next-fallback.yml
index ef8be3df..68b3e052 100644
--- a/.github/workflows/test-next-fallback.yml
+++ b/.github/workflows/test-next-fallback.yml
@@ -12,13 +12,24 @@ on:
     - "*.rst"
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.get-versions.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: get-versions
+      uses: ./.github/actions/get-python-versions
+
   test-next:
+    needs: get-python-versions
     strategy:
       matrix:
         codegen-factor: [internal, dace]
         mesh-factor: [nomesh, atlas]
         os: ["ubuntu-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
 
     runs-on: ${{ matrix.os }}
     steps:
diff --git a/.github/workflows/test-next.yml b/.github/workflows/test-next.yml
index 068377c6..68d86e5a 100644
--- a/.github/workflows/test-next.yml
+++ b/.github/workflows/test-next.yml
@@ -15,13 +15,24 @@ on:
     - "*.rst"
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.get-versions.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: get-versions
+      uses: ./.github/actions/get-python-versions
+
   test-next:
+    needs: get-python-versions
     strategy:
       matrix:
         codegen-factor: [internal, dace]
         mesh-factor: [nomesh, atlas]
         os: ["ubuntu-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
       fail-fast: false
 
     runs-on: ${{ matrix.os }}
diff --git a/.github/workflows/test-storage-fallback.yml b/.github/workflows/test-storage-fallback.yml
index c913529a..11739104 100644
--- a/.github/workflows/test-storage-fallback.yml
+++ b/.github/workflows/test-storage-fallback.yml
@@ -17,11 +17,22 @@ on:
     - "*.yml"
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.get-versions.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: get-versions
+      uses: ./.github/actions/get-python-versions
+
   test-storage:
+    needs: get-python-versions
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
 
     runs-on: ${{ matrix.os }}
     steps:
diff --git a/.github/workflows/test-storage.yml b/.github/workflows/test-storage.yml
index b2bb09df..4abedbfb 100644
--- a/.github/workflows/test-storage.yml
+++ b/.github/workflows/test-storage.yml
@@ -17,11 +17,22 @@ on:
     - "*.yml"
 
 jobs:
+  # First job to read Python versions from .python-versions file
+  get-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      python-versions: ${{ steps.get-versions.outputs.python-versions }}
+    steps:
+    - uses: actions/checkout@v4
+    - id: get-versions
+      uses: ./.github/actions/get-python-versions
+
   test-storage:
+    needs: get-python-versions
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ${{ fromJson(needs.get-python-versions.outputs.python-versions) }}
       fail-fast: false
 
     runs-on: ${{ matrix.os }}
diff --git a/tests/eve_tests/unit_tests/test_visitors.py b/tests/eve_tests/unit_tests/test_visitors.py
index 2eff41d7..c59fa167 100644
--- a/tests/eve_tests/unit_tests/test_visitors.py
+++ b/tests/eve_tests/unit_tests/test_visitors.py
@@ -24,3 +24,38 @@ def test_annex_preservation(compound_node: eve.Node):
     assert translated_node.annex.foo == 1
     assert translated_node.annex.bar is None
     assert not hasattr(translated_node.annex, "baz")
+
+
+def test_annex_preservation_translated_node(compound_node: eve.Node):
+    compound_node.annex.foo = 1
+    compound_node.annex.baz = 2
+
+    class SampleTranslator(eve.NodeTranslator):
+        PRESERVED_ANNEX_ATTRS = ("foo",)
+
+        def visit_Node(self, node: eve.Node):
+            # just return an empty node, we care about the annex only anyway
+            return eve.Node()
+
+    translated_node = SampleTranslator().visit(compound_node)
+
+    assert translated_node.annex.foo == 1
+    assert not hasattr(translated_node.annex, "baz")
+
+
+def test_annex_preservation_translated_node_overwritten(compound_node: eve.Node):
+    compound_node.annex.foo = "1+1"
+
+    class SampleTranslator(eve.NodeTranslator):
+        PRESERVED_ANNEX_ATTRS = ("foo",)
+
+        def visit_Node(self, node: eve.Node):
+            # just return an empty node, we care about the annex only anyway
+            new_node = eve.Node()
+            # the annex value is different, but considered equivalent by this pass
+            new_node.annex.foo = "2"
+            return new_node
+
+    translated_node = SampleTranslator().visit(compound_node)
+
+    assert translated_node.annex.foo == "2"