test[next]: adopt class-style dimension declarations and add pyright coverage

pyproject.toml

@@ -68,6 +68,7 @@ typing_exports = [
   {include-group = "typing"},
   'types-six>=1.17.0.20251009',  # can not let mypy auto-install types as that leads to unexpected stderr output (which means test failure)
   'pytest-mypy-plugins>=4.0.0',  # pytest plugin for running mypy on code snippets
+  'pyright>=1.1.411',  # dimension annotations must work without a gt4py plugin, see ADR 0028
   "xarray>=2024.1.0"  # one of the regression tests requires xarray
 ]
 

tests/next_tests/benchmarks/benchmark_program_call.py

@@ -41,8 +41,10 @@
     from pytest_benchmark import fixture as ptb_fixture
 
 
-Cell = gtx.dimension("Cell")
-IDim = gtx.dimension("IDim")
+class Cell(gtx.DimensionIndex): ...
+
+
+class IDim(gtx.DimensionIndex): ...
 
 
 @pytest.mark.parametrize("backend", BACKENDS, ids=lambda b: b.name)

tests/next_tests/fixtures/past_common.py

@@ -14,11 +14,14 @@
 from gt4py.next import float64
 
 
-IDim = gtx.dimension("IDim")
-JDim = gtx.dimension("JDim")
+class IDim(gtx.DimensionIndex): ...
+
+
+class JDim(
+    gtx.DimensionIndex
+): ...  # TODO(tehrengruber): Improve test structure. Identity needs to be decorated
 
 
-# TODO(tehrengruber): Improve test structure. Identity needs to be decorated
 #  in order to be used inside a program. This is unfortunate as a bug inside
 #  the decorator may result in failing tests before the actual test is run.
 #  A better way would be to first test everything field operator related,

tests/next_tests/integration_tests/cases_utils.py

@@ -153,25 +153,56 @@ def debug_itir(tree):
 DimsType = TypeVar("DimsType")
 DType = TypeVar("DType")
 
-IDim = gtx.dimension("IDim")
+
+class IDim(gtx.DimensionIndex): ...
+
+
 IHalfDim = common.flip_staggered(IDim)
-JDim = gtx.dimension("JDim")
+
+
+class JDim(gtx.DimensionIndex): ...
+
+
 JHalfDim = common.flip_staggered(JDim)
-KDim = gtx.dimension("KDim", kind=gtx.DimensionKind.VERTICAL)
+
+
+class KDim(gtx.DimensionIndex, kind=gtx.DimensionKind.VERTICAL): ...
+
+
 KHalfDim = common.flip_staggered(KDim)
 
 Ioff = gtx.FieldOffset("Ioff", source=IDim, target=(IDim,))
 Koff = gtx.FieldOffset("Koff", source=KDim, target=(KDim,))
 
-Vertex = gtx.dimension("Vertex")
-Edge = gtx.dimension("Edge")
-Cell = gtx.dimension("Cell")
+
+class Vertex(gtx.DimensionIndex): ...
+
+
+class Edge(gtx.DimensionIndex): ...
+
+
+class Cell(gtx.DimensionIndex): ...
+
+
 EdgeOffset = gtx.FieldOffset("EdgeOffset", source=Edge, target=(Edge,))
 
-V2EDim = gtx.dimension("V2E", kind=gtx.DimensionKind.LOCAL)
-E2VDim = gtx.dimension("E2V", kind=gtx.DimensionKind.LOCAL)
-C2EDim = gtx.dimension("C2E", kind=gtx.DimensionKind.LOCAL)
-C2VDim = gtx.dimension("C2V", kind=gtx.DimensionKind.LOCAL)
+
+class V2EDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "V2E"
+
+
+class E2VDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "E2V"
+
+
+class C2EDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "C2E"
+
+
+class C2VDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "C2V"
+
+
 V2E = gtx.FieldOffset("V2E", source=Edge, target=(Vertex, V2EDim))
 E2V = gtx.FieldOffset("E2V", source=Vertex, target=(Edge, E2VDim))
 C2E = gtx.FieldOffset("C2E", source=Edge, target=(Cell, C2EDim))

tests/next_tests/integration_tests/feature_tests/dace_tests/test_write_back_buffer_elimination_lowering.py

@@ -30,12 +30,23 @@
 from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations
 
 
-IDim = gtx.dimension("I")
+class IDim(gtx.DimensionIndex):
+    tag = "I"
+
+
 I_SIZE = 8
 
-Cell = gtx.dimension("Cell")
-Edge = gtx.dimension("Edge")
-C2EDim = gtx.dimension("C2E", kind=gtx.DimensionKind.LOCAL)
+
+class Cell(gtx.DimensionIndex): ...
+
+
+class Edge(gtx.DimensionIndex): ...
+
+
+class C2EDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "C2E"
+
+
 C2E = gtx.FieldOffset("C2E", source=Edge, target=(Cell, C2EDim))
 
 C2E_TABLE = np.array(

tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_hooks.py

@@ -26,7 +26,7 @@
     BACKENDS = [None, gtfn_cpu]
 
 
-IDim = gtx.dimension("IDim")
+class IDim(gtx.DimensionIndex): ...
 
 
 @gtx.field_operator

tests/next_tests/integration_tests/feature_tests/iterator_tests/test_builtins.py

@@ -67,7 +67,7 @@ def _listify(val):
     return res
 
 
-IDim = gtx.dimension("IDim")
+class IDim(gtx.DimensionIndex): ...
 
 
 def field_maker(*arrays):

tests/next_tests/integration_tests/feature_tests/iterator_tests/test_conditional.py

@@ -16,7 +16,7 @@
 from next_tests.unit_tests.conftest import program_processor, run_processor
 
 
-IDim = gtx.dimension("IDim")
+class IDim(gtx.DimensionIndex): ...
 
 
 @fundef

tests/next_tests/integration_tests/feature_tests/iterator_tests/test_implicit_fencil.py

@@ -16,7 +16,8 @@
 from next_tests.unit_tests.conftest import program_processor, run_processor
 
 
-I = gtx.dimension("I")
+class I(gtx.DimensionIndex): ...
+
 
 _isize = 10
 

tests/next_tests/integration_tests/feature_tests/iterator_tests/test_program.py

@@ -25,7 +25,9 @@
 from next_tests.unit_tests.conftest import program_processor, run_processor
 
 
-I = gtx.dimension("I")
+class I(gtx.DimensionIndex): ...
+
+
 Ioff = gtx.CartesianConnectivity(I)
 
 

tests/next_tests/integration_tests/feature_tests/iterator_tests/test_strided_offset_provider.py

@@ -17,9 +17,14 @@
 from gt4py.next.iterator.embedded import StridedConnectivityField
 
 
-LocA = gtx.dimension("LocA")
-LocAB = gtx.dimension("LocAB")
-LocB = gtx.dimension("LocB")  # unused
+class LocA(gtx.DimensionIndex): ...
+
+
+class LocAB(gtx.DimensionIndex): ...
+
+
+class LocB(gtx.DimensionIndex): ...  # unused
+
 
 LocA2LocAB = offset("O")
 LocA2LocAB_offset_provider = StridedConnectivityField(

tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py

@@ -16,11 +16,17 @@
 from next_tests.unit_tests.conftest import program_processor, run_processor
 
 
-IDim = gtx.dimension("IDim")
-JDim = gtx.dimension("JDim")
-KDim = gtx.dimension("KDim")
+class IDim(gtx.DimensionIndex): ...
+
+
+class JDim(gtx.DimensionIndex): ...
+
+
+class KDim(
+    gtx.DimensionIndex
+): ...  # semantics of stencil return that is called from the fencil (after `:` the structure of the output)
+
 
-# semantics of stencil return that is called from the fencil (after `:` the structure of the output)
 # `return a` -> a: field
 # `return make_tuple(a)` -> (a,): [field] or (field)
 # `return a,b` -> (a,b): [field, field] or (field, field)

tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py

@@ -30,8 +30,10 @@
 ]
 
 
-Cell = gtx.dimension("Cell")
-KDim = gtx.dimension("KDim", kind=gtx.DimensionKind.VERTICAL)
+class Cell(gtx.DimensionIndex): ...
+
+
+class KDim(gtx.DimensionIndex, kind=gtx.DimensionKind.VERTICAL): ...
 
 
 class State(NamedTuple):

tests/next_tests/integration_tests/multi_feature_tests/fvm_nabla_setup.py

@@ -36,10 +36,19 @@
 from gt4py.next.iterator import atlas_utils
 
 
-Vertex = gtx.dimension("Vertex")
-Edge = gtx.dimension("Edge")
-V2EDim = gtx.dimension("V2E", kind=gtx.DimensionKind.LOCAL)
-E2VDim = gtx.dimension("E2V", kind=gtx.DimensionKind.LOCAL)
+class Vertex(gtx.DimensionIndex): ...
+
+
+class Edge(gtx.DimensionIndex): ...
+
+
+class V2EDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "V2E"
+
+
+class E2VDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "E2V"
+
 
 V2E = gtx.FieldOffset("V2E", source=Edge, target=(Vertex, V2EDim))
 E2V = gtx.FieldOffset("E2V", source=Vertex, target=(Edge, E2VDim))

tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_anton_toy.py

@@ -24,12 +24,15 @@
 from next_tests.unit_tests.conftest import program_processor, run_processor
 
 
-IDim = gtx.dimension("IDim")
-JDim = gtx.dimension("JDim")
-KDim = gtx.dimension("KDim")
+class IDim(gtx.DimensionIndex): ...
+
+
+class JDim(gtx.DimensionIndex): ...
+
+
+class KDim(gtx.DimensionIndex): ...  # cross-reference why new type inference does not support this
 
 
-# cross-reference why new type inference does not support this
 @fundef
 def ldif(d):
     return lambda inp: deref(shift(d, -1)(inp)) - deref(inp)

tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_fvm_nabla.py

@@ -37,9 +37,15 @@
 from next_tests.unit_tests.conftest import program_processor, run_processor
 
 
-Vertex = gtx.dimension("Vertex")
-Edge = gtx.dimension("Edge")
-V2EDim = gtx.dimension("V2E", kind=gtx.DimensionKind.LOCAL)
+class Vertex(gtx.DimensionIndex): ...
+
+
+class Edge(gtx.DimensionIndex): ...
+
+
+class V2EDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "V2E"
+
 
 V2E = offset("V2E")
 E2V = offset("E2V")

tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_if_stmt.py

@@ -22,7 +22,7 @@ def multiply(alpha, inp):
     return deref(alpha) * deref(inp)
 
 
-IDim = gtx.dimension("IDim")
+class IDim(gtx.DimensionIndex): ...
 
 
 @pytest.mark.uses_ir_if_stmts

tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_temporaries.py

@@ -24,8 +24,11 @@
 from next_tests.unit_tests.conftest import program_processor_no_transforms, run_processor
 
 
-IDim = gtx.dimension("IDim")
-JDim = gtx.dimension("JDim")
+class IDim(gtx.DimensionIndex): ...
+
+
+class JDim(gtx.DimensionIndex): ...
+
 
 i = gtx.CartesianConnectivity(IDim)
 j = gtx.CartesianConnectivity(JDim)

tests/next_tests/regression_tests/embedded_tests/test_domain_pickle.py

@@ -10,8 +10,11 @@
 
 from gt4py.next import common
 
-I = common.dimension("I")
-J = common.dimension("J")
+
+class I(common.DimensionIndex): ...
+
+
+class J(common.DimensionIndex): ...
 
 
 def test_domain_pickle_after_slice():

tests/next_tests/regression_tests/ffront_tests/test_offset_dimensions_names.py

@@ -16,9 +16,15 @@
 from next_tests.integration_tests import cases_utils
 
 
-V = gtx.dimension("V")
-E = gtx.dimension("E")
-Neigh = gtx.dimension("Neigh", kind=common.DimensionKind.LOCAL)
+class V(gtx.DimensionIndex): ...
+
+
+class E(gtx.DimensionIndex): ...
+
+
+class Neigh(gtx.DimensionIndex, kind=common.DimensionKind.LOCAL): ...
+
+
 Off = gtx.FieldOffset("Off", source=E, target=(V, Neigh))
 
 

tests/next_tests/toy_connectivity.py

@@ -12,13 +12,30 @@
 from gt4py.next.iterator import builtins, ir as itir
 
 
-Vertex = gtx.dimension("Vertex")
-Edge = gtx.dimension("Edge")
-Cell = gtx.dimension("Cell")
-V2EDim = gtx.dimension("V2E", kind=gtx.DimensionKind.LOCAL)
-E2VDim = gtx.dimension("E2V", kind=gtx.DimensionKind.LOCAL)
-C2EDim = gtx.dimension("C2E", kind=gtx.DimensionKind.LOCAL)
-V2VDim = gtx.dimension("V2V", kind=gtx.DimensionKind.LOCAL)
+class Vertex(gtx.DimensionIndex): ...
+
+
+class Edge(gtx.DimensionIndex): ...
+
+
+class Cell(gtx.DimensionIndex): ...
+
+
+class V2EDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "V2E"
+
+
+class E2VDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "E2V"
+
+
+class C2EDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "C2E"
+
+
+class V2VDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL):
+    tag = "V2V"
+
 
 V2E = gtx.FieldOffset("V2E", source=Edge, target=(Vertex, V2EDim))
 E2V = gtx.FieldOffset("E2V", source=Vertex, target=(Edge, E2VDim))