From 9682fbf631fb7feddfd3a0a0025dbcbc98204825 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 2 Sep 2026 06:50:49 +0200 Subject: [PATCH] test[next]: migrate dimension declarations and index construction Migrates every remaining factory-style declaration to a class -- 116 under `tests/`, plus the QuickstartGuide, the workshop notebooks, `helpers.py` and the cartesian-vs-next example -- and adds type-checker coverage for the new spelling. Where the Python name differs from the dimension's tag, the class keeps the name and pins `tag` explicitly, so the backend tag and every downstream reference are unchanged. Adds pyright to the `typing_exports` group and a `test_pyright.py` runner over `pyright_cases.py` asserting diagnostics line for line. pyright has no plugin mechanism, so it is the honest check that dimension annotations work unaided; nothing exercised pyright before. The new checker cases pin what this design buys: `IDim(0)` binds to `IDim` (not a shared index type) under both checkers with no suppressions, `IDim(0).value` is `int` while `IDim.tag` is the name, passing a `JDim` index where an `IDim` is required is an `arg-type` error, and `gtx.Dimension` annotates the class so an index is not assignable to it. Notebook code cells were migrated without touching stored outputs: those hold recorded tracebacks whose text must keep naming the symbols that produced them. --- pyproject.toml | 1 + .../benchmarks/benchmark_program_call.py | 6 +- tests/next_tests/fixtures/past_common.py | 9 +- .../integration_tests/cases_utils.py | 51 +++++++-- ..._write_back_buffer_elimination_lowering.py | 19 +++- .../instrumentation_tests/test_hooks.py | 2 +- .../iterator_tests/test_builtins.py | 2 +- .../iterator_tests/test_conditional.py | 2 +- .../iterator_tests/test_implicit_fencil.py | 3 +- .../iterator_tests/test_program.py | 4 +- .../test_strided_offset_provider.py | 11 +- .../iterator_tests/test_tuple.py | 14 ++- .../ffront_tests/test_icon_like_scan.py | 6 +- .../multi_feature_tests/fvm_nabla_setup.py | 17 ++- .../iterator_tests/test_anton_toy.py | 11 +- .../iterator_tests/test_fvm_nabla.py | 12 ++- .../iterator_tests/test_if_stmt.py | 2 +- .../iterator_tests/test_temporaries.py | 7 +- .../embedded_tests/test_domain_pickle.py | 7 +- .../test_offset_dimensions_names.py | 12 ++- tests/next_tests/toy_connectivity.py | 31 ++++-- .../embedded_tests/test_basic_program.py | 2 +- .../unit_tests/embedded_tests/test_common.py | 10 +- .../test_decorator_domain_deduction.py | 7 +- .../ffront_tests/test_diagnostic_messages.py | 6 +- .../ffront_tests/test_foast_to_gtir.py | 22 +++- .../ffront_tests/test_func_to_foast.py | 4 +- .../test_func_to_foast_error_line_number.py | 3 +- .../unit_tests/ffront_tests/test_stages.py | 3 +- .../ir_utils_tests/test_domain_utils.py | 38 +++++-- .../test_embedded_field_with_list.py | 13 ++- .../test_inline_dynamic_shifts.py | 5 +- .../iterator_tests/test_runtime_domain.py | 3 +- .../transforms_tests/test_collapse_tuple.py | 4 +- ...t_concat_where_canonicalize_domain_args.py | 6 +- .../test_concat_where_expand_tuple_args.py | 6 +- ...st_concat_where_transform_to_as_fieldop.py | 8 +- .../test_dead_code_elimination.py | 5 +- .../transforms_tests/test_domain_inference.py | 27 +++-- .../test_expand_tuple_maps.py | 4 +- .../transforms_tests/test_fuse_as_fieldop.py | 8 +- .../transforms_tests/test_global_tmps.py | 12 ++- .../transforms_tests/test_inline_scalar.py | 5 +- .../test_prune_empty_concat_where.py | 17 ++- .../otf_tests/test_compiled_program.py | 2 +- .../unit_tests/test_constructors.py | 11 +- .../test_custom_layout_allocators.py | 10 +- .../test_type_translation.py | 8 +- typing_tests/pyright_cases.py | 76 +++++++++++++ typing_tests/test_next.yaml | 102 ++++++++++++++++++ typing_tests/test_pyright.py | 60 +++++++++++ uv.lock | 15 +++ 52 files changed, 610 insertions(+), 121 deletions(-) create mode 100644 typing_tests/pyright_cases.py create mode 100644 typing_tests/test_pyright.py diff --git a/pyproject.toml b/pyproject.toml index e51bac0d90..d5b281fef2 100644 --- a/pyproject.toml +++ b/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 ] diff --git a/tests/next_tests/benchmarks/benchmark_program_call.py b/tests/next_tests/benchmarks/benchmark_program_call.py index 7ecae9fb83..9f5f90b992 100644 --- a/tests/next_tests/benchmarks/benchmark_program_call.py +++ b/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) diff --git a/tests/next_tests/fixtures/past_common.py b/tests/next_tests/fixtures/past_common.py index 2430f26e0f..b38b1d0375 100644 --- a/tests/next_tests/fixtures/past_common.py +++ b/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, diff --git a/tests/next_tests/integration_tests/cases_utils.py b/tests/next_tests/integration_tests/cases_utils.py index 3426e32feb..d07b467563 100644 --- a/tests/next_tests/integration_tests/cases_utils.py +++ b/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)) diff --git a/tests/next_tests/integration_tests/feature_tests/dace_tests/test_write_back_buffer_elimination_lowering.py b/tests/next_tests/integration_tests/feature_tests/dace_tests/test_write_back_buffer_elimination_lowering.py index 928322cc71..d4421f9b28 100644 --- a/tests/next_tests/integration_tests/feature_tests/dace_tests/test_write_back_buffer_elimination_lowering.py +++ b/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( diff --git a/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_hooks.py b/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_hooks.py index 02c51a8557..3700eac616 100644 --- a/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_hooks.py +++ b/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 diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_builtins.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_builtins.py index f6fb3852da..57e8d9f254 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_builtins.py +++ b/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): diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_conditional.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_conditional.py index 657f3c7596..ed4039567d 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_conditional.py +++ b/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 diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_implicit_fencil.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_implicit_fencil.py index 2dccd6dcec..ed1ed9d6a8 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_implicit_fencil.py +++ b/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 diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_program.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_program.py index 72d9868f63..7220830aab 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_program.py +++ b/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) diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_strided_offset_provider.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_strided_offset_provider.py index 482db69614..3f02a1a359 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_strided_offset_provider.py +++ b/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( diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py index f152e69e3c..d18787d422 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py +++ b/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) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py index ad579f90bf..5b40426a0e 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py +++ b/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): diff --git a/tests/next_tests/integration_tests/multi_feature_tests/fvm_nabla_setup.py b/tests/next_tests/integration_tests/multi_feature_tests/fvm_nabla_setup.py index f4cca4cb19..1257760160 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/fvm_nabla_setup.py +++ b/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)) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_anton_toy.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_anton_toy.py index 7bfe4da0f7..7e09c577d8 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_anton_toy.py +++ b/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) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_fvm_nabla.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_fvm_nabla.py index 910cb33145..3e337ca407 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_fvm_nabla.py +++ b/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") diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_if_stmt.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_if_stmt.py index 56aa8e1c2c..746e9f83ed 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_if_stmt.py +++ b/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 diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_temporaries.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_temporaries.py index de086fd24a..6db3cafd2d 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_temporaries.py +++ b/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) diff --git a/tests/next_tests/regression_tests/embedded_tests/test_domain_pickle.py b/tests/next_tests/regression_tests/embedded_tests/test_domain_pickle.py index 6a6af63cbf..b5588ff7ed 100644 --- a/tests/next_tests/regression_tests/embedded_tests/test_domain_pickle.py +++ b/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(): diff --git a/tests/next_tests/regression_tests/ffront_tests/test_offset_dimensions_names.py b/tests/next_tests/regression_tests/ffront_tests/test_offset_dimensions_names.py index bd9e42c88f..a12d070a4e 100644 --- a/tests/next_tests/regression_tests/ffront_tests/test_offset_dimensions_names.py +++ b/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)) diff --git a/tests/next_tests/toy_connectivity.py b/tests/next_tests/toy_connectivity.py index 3e9118a950..88fc4e9491 100644 --- a/tests/next_tests/toy_connectivity.py +++ b/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)) diff --git a/tests/next_tests/unit_tests/embedded_tests/test_basic_program.py b/tests/next_tests/unit_tests/embedded_tests/test_basic_program.py index cde2ca4f27..f5a3212889 100644 --- a/tests/next_tests/unit_tests/embedded_tests/test_basic_program.py +++ b/tests/next_tests/unit_tests/embedded_tests/test_basic_program.py @@ -11,7 +11,7 @@ import gt4py.next as gtx -IDim = gtx.dimension("IDim") +class IDim(gtx.DimensionIndex): ... @gtx.field_operator diff --git a/tests/next_tests/unit_tests/embedded_tests/test_common.py b/tests/next_tests/unit_tests/embedded_tests/test_common.py index df733e93b9..2cf2cf3b47 100644 --- a/tests/next_tests/unit_tests/embedded_tests/test_common.py +++ b/tests/next_tests/unit_tests/embedded_tests/test_common.py @@ -37,9 +37,13 @@ def test_slice_range(rng, slce, expected): assert result == expected -I = common.dimension("I") -J = common.dimension("J") -K = common.dimension("K") +class I(common.DimensionIndex): ... + + +class J(common.DimensionIndex): ... + + +class K(common.DimensionIndex): ... @pytest.mark.parametrize( diff --git a/tests/next_tests/unit_tests/ffront_tests/test_decorator_domain_deduction.py b/tests/next_tests/unit_tests/ffront_tests/test_decorator_domain_deduction.py index ec80de981e..9bee36576a 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_decorator_domain_deduction.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_decorator_domain_deduction.py @@ -12,8 +12,11 @@ from gt4py.next.ffront.transform_utils import _deduce_grid_type -Dim = gtx.dimension("Dim") -LocalDim = gtx.dimension("LocalDim", kind=gtx.DimensionKind.LOCAL) +class Dim(gtx.DimensionIndex): ... + + +class LocalDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL): ... + CartesianOffset = gtx.FieldOffset("CartesianOffset", source=Dim, target=(Dim,)) UnstructuredOffset = gtx.FieldOffset("UnstructuredOffset", source=Dim, target=(Dim, LocalDim)) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py index 6c442062fa..2f80faa0dd 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py @@ -27,9 +27,11 @@ from gt4py.next.ffront.func_to_foast import FieldOperatorParser -IDim = gtx.dimension("IDim") +class IDim( + gtx.DimensionIndex +): ... # A PEP 695 alias whose value raises when it is evaluated, standing in for the + -# A PEP 695 alias whose value raises when it is evaluated, standing in for the # common case of a typo'd dtype ('np.foat64') inside an alias definition. _empty_module = types.ModuleType("_empty_module") type BrokenFieldAlias = gtx.Field[gtx.Dims[IDim], _empty_module.foat64] diff --git a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py index 47efeec1d6..f4d4f6c8cd 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py @@ -43,14 +43,26 @@ from gt4py.next.iterator import ir as itir -Edge = gtx.dimension("Edge") -Vertex = gtx.dimension("Vertex") -V2EDim = gtx.dimension("V2E", gtx.DimensionKind.LOCAL) +class Edge(gtx.DimensionIndex): ... + + +class Vertex(gtx.DimensionIndex): ... + + +class V2EDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL): + tag = "V2E" + + V2E = gtx.FieldOffset("V2E", source=Edge, target=(Vertex, V2EDim)) -TDim = gtx.dimension("TDim") + +class TDim(gtx.DimensionIndex): ... + + TOff = gtx.FieldOffset("TDim", source=TDim, target=(TDim,)) -UDim = gtx.dimension("UDim") + + +class UDim(gtx.DimensionIndex): ... def test_return(): diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py index a2c29498c8..1504d3a064 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py @@ -71,7 +71,9 @@ XOR = itir.SymRef(id=itb.xor_.fun.__name__) LIFT = itir.SymRef(id=itb.lift.fun.__name__) -TDim = gtx.dimension("TDim") # Meaningless dimension, used for tests. + +class TDim(gtx.DimensionIndex): ... # Meaningless dimension, used for tests. + # PEP 695 type alias, used to check that aliases are accepted as DSL annotations. type TFloatFieldAlias = gtx.Field[gtx.Dims[TDim], float64] diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast_error_line_number.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast_error_line_number.py index c9117af3bf..1a98d4c67f 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast_error_line_number.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast_error_line_number.py @@ -19,7 +19,8 @@ # NOTE: These tests are sensitive to filename and the line number of the marked statement -TDim = gtx.dimension("TDim") # Meaningless dimension, used for tests. + +class TDim(gtx.DimensionIndex): ... # Meaningless dimension, used for tests. def test_invalid_syntax_error_empty_return(): diff --git a/tests/next_tests/unit_tests/ffront_tests/test_stages.py b/tests/next_tests/unit_tests/ffront_tests/test_stages.py index fa6f92b45f..ab02fac429 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_stages.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_stages.py @@ -12,7 +12,8 @@ from gt4py.next.ffront import stages -IDim = gtx.dimension("I") +class IDim(gtx.DimensionIndex): + tag = "I" def _make_field_operator_definition(offset: int): diff --git a/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py b/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py index 6022ed587a..6262741161 100644 --- a/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py +++ b/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py @@ -14,15 +14,37 @@ from gt4py.next.iterator.ir_utils import domain_utils, ir_makers as im from gt4py.next import common, constructors -I = common.dimension("I") + +class I(common.DimensionIndex): ... + + IHalf = common.flip_staggered(I) -J = common.dimension("J") -K = common.dimension("J", kind=common.DimensionKind.VERTICAL) -Vertex = common.dimension("Vertex") -Edge = common.dimension("Edge") -V2EDim = common.dimension("V2E", kind=common.DimensionKind.LOCAL) -E2VDim = common.dimension("E2V", kind=common.DimensionKind.LOCAL) -V2VDim = common.dimension("V2V", kind=common.DimensionKind.LOCAL) + + +class J(common.DimensionIndex): ... + + +class K(common.DimensionIndex, kind=common.DimensionKind.VERTICAL): + tag = "J" + + +class Vertex(common.DimensionIndex): ... + + +class Edge(common.DimensionIndex): ... + + +class V2EDim(common.DimensionIndex, kind=common.DimensionKind.LOCAL): + tag = "V2E" + + +class E2VDim(common.DimensionIndex, kind=common.DimensionKind.LOCAL): + tag = "E2V" + + +class V2VDim(common.DimensionIndex, kind=common.DimensionKind.LOCAL): + tag = "V2V" + a_range = domain_utils.SymbolicRange(0, 10) another_range = domain_utils.SymbolicRange(5, 15) diff --git a/tests/next_tests/unit_tests/iterator_tests/test_embedded_field_with_list.py b/tests/next_tests/unit_tests/iterator_tests/test_embedded_field_with_list.py index 6e45e663a4..4db1f80767 100644 --- a/tests/next_tests/unit_tests/iterator_tests/test_embedded_field_with_list.py +++ b/tests/next_tests/unit_tests/iterator_tests/test_embedded_field_with_list.py @@ -23,9 +23,16 @@ ) -E = gtx.dimension("E") -V = gtx.dimension("V") -E2VDim = gtx.dimension("E2V", kind=gtx.DimensionKind.LOCAL) +class E(gtx.DimensionIndex): ... + + +class V(gtx.DimensionIndex): ... + + +class E2VDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL): + tag = "E2V" + + E2V = gtx.FieldOffset("E2V", source=V, target=(E, E2VDim)) diff --git a/tests/next_tests/unit_tests/iterator_tests/test_inline_dynamic_shifts.py b/tests/next_tests/unit_tests/iterator_tests/test_inline_dynamic_shifts.py index e4a5e0b862..b5432098bc 100644 --- a/tests/next_tests/unit_tests/iterator_tests/test_inline_dynamic_shifts.py +++ b/tests/next_tests/unit_tests/iterator_tests/test_inline_dynamic_shifts.py @@ -11,7 +11,10 @@ from gt4py.next.iterator.transforms import inline_dynamic_shifts from gt4py.next.type_system import type_specifications as ts -IDim = gtx.dimension("IDim") + +class IDim(gtx.DimensionIndex): ... + + field_type = ts.FieldType(dims=[IDim], dtype=ts.ScalarType(kind=ts.ScalarKind.INT32)) IOff = im.cartesian_offset(IDim, IDim) diff --git a/tests/next_tests/unit_tests/iterator_tests/test_runtime_domain.py b/tests/next_tests/unit_tests/iterator_tests/test_runtime_domain.py index c98e3c8292..443eadfaf1 100644 --- a/tests/next_tests/unit_tests/iterator_tests/test_runtime_domain.py +++ b/tests/next_tests/unit_tests/iterator_tests/test_runtime_domain.py @@ -27,7 +27,8 @@ def foo(inp): dtype=None, ) -I = gtx.dimension("I") + +class I(gtx.DimensionIndex): ... def test_deduce_domain(): diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py index 39e808c87b..f69fbc768e 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py @@ -15,7 +15,9 @@ int_type = ts.ScalarType(kind=ts.ScalarKind.INT32) -Vertex = common.dimension("Vertex", kind=common.DimensionKind.HORIZONTAL) + + +class Vertex(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... def test_simple_make_tuple_tuple_get(uids: utils.IDGeneratorPool): diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_canonicalize_domain_args.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_canonicalize_domain_args.py index 5b64eb519f..b6a7d05388 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_canonicalize_domain_args.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_canonicalize_domain_args.py @@ -16,7 +16,11 @@ from gt4py.next.type_system import type_specifications as ts int_type = ts.ScalarType(kind=ts.ScalarKind.INT32) -IDim = common.dimension("IDim", kind=common.DimensionKind.HORIZONTAL) + + +class IDim(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + field_type = ts.FieldType(dims=[IDim], dtype=int_type) diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_expand_tuple_args.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_expand_tuple_args.py index d341ca8ffe..97c2f4ce40 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_expand_tuple_args.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_expand_tuple_args.py @@ -17,7 +17,11 @@ from gt4py.next.iterator.type_system import type_specifications as it_ts int_type = ts.ScalarType(kind=ts.ScalarKind.INT32) -IDim = common.dimension("IDim", kind=common.DimensionKind.HORIZONTAL) + + +class IDim(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + field_type = ts.FieldType(dims=[IDim], dtype=int_type) diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_transform_to_as_fieldop.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_transform_to_as_fieldop.py index 203d587e11..babcafc4a1 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_transform_to_as_fieldop.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_transform_to_as_fieldop.py @@ -14,8 +14,12 @@ from gt4py.next.type_system import type_specifications as ts int_type = ts.ScalarType(kind=ts.ScalarKind.INT32) -IDim = common.dimension("IDim", kind=common.DimensionKind.HORIZONTAL) -JDim = common.dimension("JDim", kind=common.DimensionKind.HORIZONTAL) + + +class IDim(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + +class JDim(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... def test_in_helper(): diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_dead_code_elimination.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_dead_code_elimination.py index 85dbdc8d14..f2e4f4bb07 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_dead_code_elimination.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_dead_code_elimination.py @@ -14,7 +14,10 @@ from gt4py.next.iterator import ir as itir from gt4py.next.iterator.transforms import dead_code_elimination -TDim = common.dimension("TDim") + +class TDim(common.DimensionIndex): ... + + int_type = ts.ScalarType(kind=ts.ScalarKind.INT32) field_type = ts.FieldType(dims=[TDim], dtype=int_type) diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py index 571d6ed464..5f68a87484 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py @@ -28,12 +28,27 @@ float_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) -IDim = common.dimension("IDim", kind=common.DimensionKind.HORIZONTAL) -JDim = common.dimension("JDim", kind=common.DimensionKind.HORIZONTAL) -KDim = common.dimension("KDim", kind=common.DimensionKind.VERTICAL) -Vertex = common.dimension("Vertex", kind=common.DimensionKind.HORIZONTAL) -Edge = common.dimension("Edge", kind=common.DimensionKind.HORIZONTAL) -E2VDim = common.dimension("E2V", kind=common.DimensionKind.LOCAL) + + +class IDim(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + +class JDim(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + +class KDim(common.DimensionIndex, kind=common.DimensionKind.VERTICAL): ... + + +class Vertex(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + +class Edge(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + +class E2VDim(common.DimensionIndex, kind=common.DimensionKind.LOCAL): + tag = "E2V" + + float_i_field = ts.FieldType(dims=[IDim], dtype=float_type) float_ij_field = ts.FieldType(dims=[IDim, JDim], dtype=float_type) tuple_float_i_field = ts.TupleType( diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_expand_tuple_maps.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_expand_tuple_maps.py index f4464c7152..d8f3327386 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_expand_tuple_maps.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_expand_tuple_maps.py @@ -16,7 +16,9 @@ from gt4py.next.type_system import type_specifications as ts -IDim = common.dimension("IDim") +class IDim(common.DimensionIndex): ... + + T = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) i_field = ts.FieldType(dims=[IDim], dtype=T) i_tuple_field = ts.TupleType(types=[i_field, i_field]) diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_fuse_as_fieldop.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_fuse_as_fieldop.py index 168a0e8b4b..4177bf5681 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_fuse_as_fieldop.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_fuse_as_fieldop.py @@ -17,8 +17,12 @@ from gt4py.next.type_system import type_specifications as ts -IDim = common.dimension("IDim") -JDim = common.dimension("JDim") +class IDim(common.DimensionIndex): ... + + +class JDim(common.DimensionIndex): ... + + field_type = ts.FieldType(dims=[IDim], dtype=ts.ScalarType(kind=ts.ScalarKind.INT32)) IOff = im.cartesian_offset(IDim, IDim) diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py index bddf145d97..9adfdaedbd 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py @@ -26,9 +26,15 @@ ) -IDim = common.dimension("IDim") -JDim = common.dimension("JDim") -KDim = common.dimension("KDim", kind=common.DimensionKind.VERTICAL) +class IDim(common.DimensionIndex): ... + + +class JDim(common.DimensionIndex): ... + + +class KDim(common.DimensionIndex, kind=common.DimensionKind.VERTICAL): ... + + index_type = ts.ScalarType(kind=getattr(ts.ScalarKind, builtins.INTEGER_INDEX_BUILTIN.upper())) float_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) i_field_type = ts.FieldType(dims=[IDim], dtype=float_type) diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_scalar.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_scalar.py index fdfb2cb512..3a7c92b784 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_scalar.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_scalar.py @@ -13,7 +13,10 @@ from gt4py.next.iterator.transforms import inline_scalar from gt4py.next.iterator.ir_utils import ir_makers as im -TDim = common.dimension("TDim") + +class TDim(common.DimensionIndex): ... + + int_type = ts.ScalarType(kind=ts.ScalarKind.INT32) diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_prune_empty_concat_where.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_prune_empty_concat_where.py index 9d1a35ab21..eb9e8394c8 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_prune_empty_concat_where.py +++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_prune_empty_concat_where.py @@ -18,10 +18,19 @@ from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, domain_utils from gt4py.next.type_system import type_info, type_specifications as ts -Vertex = common.dimension("Vertex", kind=common.DimensionKind.HORIZONTAL) -Edge = common.dimension("Edge", kind=common.DimensionKind.HORIZONTAL) -V2EDim = common.dimension("V2E", kind=common.DimensionKind.LOCAL) -K = common.dimension("K", kind=common.DimensionKind.VERTICAL) + +class Vertex(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + +class Edge(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + +class V2EDim(common.DimensionIndex, kind=common.DimensionKind.LOCAL): + tag = "V2E" + + +class K(common.DimensionIndex, kind=common.DimensionKind.VERTICAL): ... + float64 = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) vertex_k_field = ts.FieldType(dims=[Vertex, K], dtype=float64) diff --git a/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py b/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py index 9aa76f78a8..45c1241dbe 100644 --- a/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py +++ b/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py @@ -62,7 +62,7 @@ def test_sanitize_static_args_wrong_type(): static_arg.validate("foo", ts.ScalarType(kind=ts.ScalarKind.INT32)) -TDim = gtx.dimension("TDim") +class TDim(gtx.DimensionIndex): ... @pytest.fixture diff --git a/tests/next_tests/unit_tests/test_constructors.py b/tests/next_tests/unit_tests/test_constructors.py index a2809d8fa6..3bfdacbf63 100644 --- a/tests/next_tests/unit_tests/test_constructors.py +++ b/tests/next_tests/unit_tests/test_constructors.py @@ -22,9 +22,14 @@ ) -I = gtx.dimension("I") -J = gtx.dimension("J") -K = gtx.dimension("K") +class I(gtx.DimensionIndex): ... + + +class J(gtx.DimensionIndex): ... + + +class K(gtx.DimensionIndex): ... + sizes = {I: 10, J: 10, K: 10} diff --git a/tests/next_tests/unit_tests/test_custom_layout_allocators.py b/tests/next_tests/unit_tests/test_custom_layout_allocators.py index 5b849f20f5..b4b5055d15 100644 --- a/tests/next_tests/unit_tests/test_custom_layout_allocators.py +++ b/tests/next_tests/unit_tests/test_custom_layout_allocators.py @@ -124,9 +124,13 @@ def test_horizontal_first_layout_mapper(): assert horizontal_first_layout_mapper(dims) == expected_layout_map -Cell = common.dimension("Cell", common.DimensionKind.HORIZONTAL) -Edge = common.dimension("Edge", common.DimensionKind.HORIZONTAL) -K = common.dimension("K", common.DimensionKind.VERTICAL) +class Cell(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + +class Edge(common.DimensionIndex, kind=common.DimensionKind.HORIZONTAL): ... + + +class K(common.DimensionIndex, kind=common.DimensionKind.VERTICAL): ... class TestBaseFieldBufferAllocatorAlignedIndex: diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py index 5da19be60e..986e3660f5 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py @@ -30,10 +30,12 @@ def dtype(self) -> np.dtype: return np.dtype(np.int32) -IDim = gtx.dimension("IDim") -JDim = gtx.dimension("JDim") +class IDim(gtx.DimensionIndex): ... + + +class JDim(gtx.DimensionIndex): ... # -- PEP 695 type aliases -- + -# -- PEP 695 type aliases -- type IFloatFieldAlias = gtx.Field[gtx.Dims[IDim], float] type ChainedFieldAlias = IFloatFieldAlias type GenericFieldAlias[T] = gtx.Field[gtx.Dims[IDim], T] diff --git a/typing_tests/pyright_cases.py b/typing_tests/pyright_cases.py new file mode 100644 index 0000000000..b73ca68c5f --- /dev/null +++ b/typing_tests/pyright_cases.py @@ -0,0 +1,76 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +""" +Cases checked by pyright in 'test_pyright.py'. + +pyright must accept every construct here with zero diagnostics and with no gt4py plugin (pyright +has no plugin mechanism, which is exactly the point -- see ADR 0028). Lines that are *expected* to +error carry a trailing EXPECT-ERROR marker comment, which the test scans for and asserts on. +""" + +from typing import TypeVar, assert_type + +from gt4py import next as gtx +from gt4py.next import common + + +class IDim(gtx.DimensionIndex): ... + + +class JDim(gtx.DimensionIndex): ... + + +class KDim(gtx.DimensionIndex, kind=gtx.DimensionKind.VERTICAL): ... + + +# 1. A dimension class is a valid annotation argument. +a: gtx.Field[gtx.Dims[IDim, KDim], gtx.float64] + +# 2. Value-level use of the class object: the dimension's name is `tag`. +assert_type(IDim.tag, common.Tag) + +# 3. An index is an ordinary instance of its dimension -- no metaclass `__call__` overloads and +# no `type: ignore` are needed to say so, and mypy and pyright agree natively. +assert_type(IDim(0), IDim) +assert_type(IDim(0).value, int) +assert_type(IDim(0).dim, gtx.Dimension) + + +# 4. Indices carry their dimension in the type, which a shared index type could not express. +def takes_i_index(i: IDim) -> None: ... + + +def pass_j_index(j: JDim) -> None: + takes_i_index(j) # EXPECT-ERROR + + +# 5. TypeVars can range over dimensions -- impossible with the old mypy plugin. +D = TypeVar("D", bound=gtx.DimensionIndex) + + +def name_of(d: type[D]) -> common.Tag: + return d.tag + + +assert_type(name_of(IDim), common.Tag) + + +# 6. A dimension mismatch between fields is a static error. +def takes_i_field(f: gtx.Field[gtx.Dims[IDim], gtx.float64]) -> None: ... + + +def pass_j_field(g: gtx.Field[gtx.Dims[JDim], gtx.float64]) -> None: + takes_i_field(g) # EXPECT-ERROR + + +# 7. `gtx.Dimension` annotates the class object, so an *index* is not one of them. +ok: gtx.Dimension = IDim +bad: gtx.Dimension = IDim(0) # EXPECT-ERROR + +# 8. The programmatic constructor returns a dimension, not an index. +assert_type(gtx.dimension("Runtime"), gtx.Dimension) diff --git a/typing_tests/test_next.yaml b/typing_tests/test_next.yaml index 20ed020333..fa295c2a65 100644 --- a/typing_tests/test_next.yaml +++ b/typing_tests/test_next.yaml @@ -275,3 +275,105 @@ main: | import xarray a: xarray.NamedArray + + - case: class_style_dimension_in_field_annotation + main: | + from gt4py import next as gtx + + class IDim(gtx.DimensionIndex): ... + class KDim(gtx.DimensionIndex, kind=gtx.DimensionKind.VERTICAL): ... + + a: gtx.Field[gtx.Dims[IDim, KDim], gtx.float64] + reveal_type(a) + out: | + main:7:13: note: Revealed type is "gt4py.next.common.Field[tuple[main.IDim, main.KDim, fallback=gt4py.next.common.Dims[main.IDim, main.KDim]], float]" + + - case: class_style_dimension_mismatch_is_an_error + main: | + from gt4py import next as gtx + + class IDim(gtx.DimensionIndex): ... + class JDim(gtx.DimensionIndex): ... + + def takes_i(f: gtx.Field[gtx.Dims[IDim], gtx.float64]) -> None: ... + def pass_j(g: gtx.Field[gtx.Dims[JDim], gtx.float64]) -> None: + takes_i(g) + out: | + main:8:13: error: Argument 1 to "takes_i" has incompatible type "Field[Dims[JDim], float]"; expected "Field[Dims[IDim], float]" [arg-type] + main:8:13: note: "Field[Dims[IDim], float].__call__" has type "def __call__(self, index_field: Connectivity[Any, Any] | FieldOffset, *args: Connectivity[Any, Any] | FieldOffset) -> Field[Any, Any]" + + - case: dimension_typevar_bound + main: | + from typing import TypeVar + from gt4py import next as gtx + + D = TypeVar("D", bound=gtx.DimensionIndex) + + class IDim(gtx.DimensionIndex): ... + + def name_of(d: type[D]) -> str: + return d.tag + + reveal_type(name_of(IDim)) + out: | + main:11:13: note: Revealed type is "str" + + - case: index_is_an_instance_of_its_dimension + main: | + from gt4py import next as gtx + + class IDim(gtx.DimensionIndex): ... + + # no expected output: `IDim(0)` is an ordinary instance of `IDim`, and its position is + # `.value`; the dimension's name is `IDim.tag`. + idx: IDim = IDim(0) + pos: int = IDim(0).value + dim: gtx.Dimension = IDim(0).dim + name: str = IDim.tag + + - case: index_of_the_wrong_dimension_is_an_error + main: | + from gt4py import next as gtx + + class IDim(gtx.DimensionIndex): ... + class JDim(gtx.DimensionIndex): ... + + def takes_i(i: IDim) -> None: ... + + takes_i(JDim(0)) + out: | + main:8:9: error: Argument 1 to "takes_i" has incompatible type "JDim"; expected "IDim" [arg-type] + + - case: unparameterized_dims_still_accepts_any_rank + main: | + from gt4py import next as gtx + + def anything(f: gtx.Field[gtx.Dims, gtx.float64]) -> None: ... + + class IDim(gtx.DimensionIndex): ... + class JDim(gtx.DimensionIndex): ... + + a: gtx.Field[gtx.Dims[IDim], gtx.float64] + b: gtx.Field[gtx.Dims[IDim, JDim], gtx.float64] + anything(a) + anything(b) + + - case: programmatic_constructor_returns_a_dimension + main: | + from gt4py import next as gtx + + IDim = gtx.dimension("IDim") + reveal_type(IDim) + out: | + main:4:13: note: Revealed type is "type[gt4py.next.common.DimensionIndex]" + + - case: dimension_alias_annotates_the_class_not_an_index + main: | + from gt4py import next as gtx + + class IDim(gtx.DimensionIndex): ... + + ok: gtx.Dimension = IDim + bad: gtx.Dimension = IDim(0) + out: | + main:6:22: error: Incompatible types in assignment (expression has type "IDim", variable has type "type[DimensionIndex]") [assignment] diff --git a/typing_tests/test_pyright.py b/typing_tests/test_pyright.py new file mode 100644 index 0000000000..7caca121f6 --- /dev/null +++ b/typing_tests/test_pyright.py @@ -0,0 +1,60 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +""" +Run pyright over 'pyright_cases.py' and assert its diagnostics. + +This exists because the dimension half of the mypy plugin was removed (ADR 0028): client code must +now type-check under a checker that has never had gt4py support. +""" + +import json +import pathlib +import shutil +import subprocess +import sys + +import pytest + + +CASES = pathlib.Path(__file__).parent / "pyright_cases.py" +MARKER = "# EXPECT-ERROR" + + +def _expected_error_lines() -> set[int]: + return {n for n, line in enumerate(CASES.read_text().splitlines(), start=1) if MARKER in line} + + +@pytest.mark.skipif(shutil.which("pyright") is None, reason="pyright is not installed") +def test_pyright_diagnostics_match_expectations(): + result = subprocess.run( + [ + "pyright", + "--outputjson", + "--pythonversion", + f"{sys.version_info.major}.{sys.version_info.minor}", + str(CASES), + ], + capture_output=True, + text=True, + ) + report = json.loads(result.stdout) + errors = { + d["range"]["start"]["line"] + 1 + for d in report["generalDiagnostics"] + if d["severity"] == "error" + } + expected = _expected_error_lines() + + assert errors == expected, ( + f"unexpected pyright errors on lines {sorted(errors - expected)}; " + f"missing expected errors on lines {sorted(expected - errors)}\n" + + "\n".join( + f" {d['range']['start']['line'] + 1}: {d['message'].splitlines()[0]}" + for d in report["generalDiagnostics"] + ) + ) diff --git a/uv.lock b/uv.lock index 3b81092dab..7303830af9 100644 --- a/uv.lock +++ b/uv.lock @@ -1513,6 +1513,7 @@ typing = [ ] typing-exports = [ { name = "mypy", extra = ["faster-cache"] }, + { name = "pyright" }, { name = "pytest-mypy-plugins" }, { name = "types-decorator" }, { name = "types-docutils" }, @@ -1672,6 +1673,7 @@ typing = [ ] typing-exports = [ { name = "mypy", extras = ["faster-cache"], specifier = ">=1.13.0" }, + { name = "pyright", specifier = ">=1.1.411" }, { name = "pytest-mypy-plugins", specifier = ">=4.0.0" }, { name = "types-decorator", specifier = ">=5.1.8" }, { name = "types-docutils", specifier = ">=0.21.0" }, @@ -3200,6 +3202,19 @@ version = "2.1" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/bc/7c/d724ef1ec3ab2125f38a1d53285745445ec4a8f19b9bb0761b4064316679/pyreadline-2.1.zip", hash = "sha256:4530592fc2e85b25b1a9f79664433da09237c1a270e4d78ea5aa3a2c7229e2d1", size = 109189, upload-time = "2015-09-16T08:24:48.745Z" } +[[package]] +name = "pyright" +version = "1.1.411" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nodeenv" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7e/ab/265f7dc69d28113ebba19092e57b075f41543b2ed048429c5f56e2b88eac/pyright-1.1.411.tar.gz", hash = "sha256:d885a0551f2e763b089a02702174e7f4ba77548cddabc972ab86d1f7f1b0f998", size = 4112861, upload-time = "2026-06-25T02:14:06.37Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0a/49/385be530a6a5b78d1cbcd5c2e38debc8959a2fc6bdb716f4e581002979fc/pyright-1.1.411-py3-none-any.whl", hash = "sha256:dc7c72a8e2700c55baa127554040e067041ea53ccfd50bf96308cc4291c7d5d9", size = 6181526, upload-time = "2026-06-25T02:14:04.691Z" }, +] + [[package]] name = "pytest" version = "9.1.1"