fix[next]: Reduce type ignores in client code

[DropD]

Description

Partner-PR in icon4py: C2SM/icon4py#1096 (make sure this passes tests first)

Radically reduce the amount of required type ignores in client code which uses the from gt4py import next as gtx pattern.

• fixes next: type hints for dsl decorators need to be improved. #2496
• fixes next: dimension types are not valid types for static type checking purposes #2503

Experimenting with this in icon4py reveals more ways in which gt4py blocks client code type checking, which have not been addressed in this PR so far:

• next: builtin function type hints block client type checking #2507
• subtly different versions of next: builtin function type hints block client type checking #2507 where mypy reports a mismatch in dtypes, which have not been looked into further so far

Requirements

• All fixes and/or new features come with corresponding tests.
• Important design decisions have been documented in the appropriate ADR inside the docs/development/ADRs/ folder.

[egparedes]

Thanks for the PR. I just have a comment and a question

[DropD]

cscs-ci run default

[DropD]

cscs-ci run default

[Copilot]

Pull request overview

This PR improves the static typing experience for client code using from gt4py import next as gtx, aiming to reduce type: ignore usage and add CI coverage for “typed client usability” scenarios.

Changes:

• Add a mypy plugin to treat Dims[...] and *Dim symbols as valid type-checking types and to blur scalar dtype precision in some cases.
• Improve exported typing surface by refining decorator signatures (program, field_operator, scan_operator) and enhancing builtin typings (notably where, astype) plus explicit fbuiltins exports.
• Add a dedicated typing test suite (pytest-mypy-plugins) and wire it into nox + GitHub Actions.

Reviewed changes

Copilot reviewed 14 out of 15 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
uv.lock	Updates locked dependencies and adds packages needed for typing-export checks (e.g., pytest-mypy-plugins, xarray).
typing_tests/test_next.yaml	Adds mypy-based “client usability” typing tests for decorators, dims, builtins, and exports.
src/gt4py/next/type_system/mypy_plugin.py	Introduces a mypy plugin to treat dimensions/dims as types and blur dtype precision.
src/gt4py/next/program_processors/runners/dace/transformations/remove_scalar_copies.py	Fixes docstring escaping and a spelling error.
src/gt4py/next/program_processors/runners/dace/transformations/loop_blocking.py	Fixes invalid escape sequences in docstring LaTeX.
src/gt4py/next/iterator/embedded.py	Removes some type ignores and adds comparison operator stubs for typing.
src/gt4py/next/ffront/fbuiltins.py	Makes exports explicit, adds overloads for builtins, and restructures math builtin registration.
src/gt4py/next/ffront/decorator.py	Broadens decorator typing from FunctionType to Callable and adjusts overload defaults.
src/gt4py/next/embedded/nd_array_field.py	Removes attr-defined ignores for builtins; adjusts typing around overloaded astype.
src/gt4py/next/common.py	Adds TYPE_CHECKING-only placeholder dim classes and expands Field protocol with comparison operators.
src/gt4py/next/__init__.py	Replaces star re-exports with explicit imports and an explicit __all__.
src/gt4py/_core/definitions.py	Adds a targeted mypy ignore related to new dtype-precision blurring behavior.
pyproject.toml	Enables the new mypy plugin and adds a dependency group for typing-export tests.
noxfile.py	Adds a nox session to run typing-export checks via pytest-mypy-plugins.
.github/workflows/code-quality.yml	Adds a CI job to run the typing-export checks.

Comments suppressed due to low confidence (2)

.github/workflows/code-quality.yml:64

• fromJson(...) is not a valid GitHub Actions expression function name (the built-in is fromJSON). As written, this step will fail to render the session name and the job will error before running nox; use the correct function name (or reuse the existing fromJSON pattern used earlier in this workflow).

      - name: "Run typing checks on exported entities"
        run: |
          ./noxfile.py -s 'test_typing_exports-${{ fromJson(needs.get-python-versions.outputs.python-versions) }}'

src/gt4py/next/ffront/fbuiltins.py:246

• The where overload that returns Tuple accepts true_field: Tuple | Scalar and false_field: Tuple | Scalar, but the runtime implementation explicitly raises ValueError when only one side is a tuple. This makes the public typing contract disagree with actual behavior (and can let mypy accept code that will fail at runtime). Either narrow the tuple overload to require tuples on both sides, or extend the implementation to support broadcasting a scalar to a tuple (and update the error checks accordingly).

    @overload  # type: ignore[override] # this technically clashes with the superclass
    def __call__(  # type: ignore[overload-overlap] # without both overloads mypy raises false positives
        self,
        cond: CondT,
        true_field: common.Field | core_defs.ScalarT,
        false_field: common.Field | core_defs.ScalarT,
    ) -> common.Field: ...

    @overload
    def __call__(
        self,
        cond: CondT,
        true_field: Tuple | core_defs.ScalarT,
        false_field: Tuple | core_defs.ScalarT,
    ) -> Tuple: ...

    def __call__(self, cond: CondT, true_field: FieldT1, false_field: FieldT2) -> _R:  # type: ignore[misc] # supposedly this signature does not accept all the possible args allowed by the overloads ??
        if isinstance(true_field, tuple) or isinstance(false_field, tuple):
            if not (isinstance(true_field, tuple) and isinstance(false_field, tuple)):
                raise ValueError(
                    # TODO(havogt) find a strategy to unify parsing and embedded error messages
                    f"Either both or none can be tuple in '{true_field=}' and '{false_field=}'."
                )
            if len(true_field) != len(false_field):

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

[egparedes]

Mostly questions. The only important issue remaining is the lack of the documentation in the mypy plugin. Other than that, it looks good to me. I have also requested a review from Copilot in case it finds something useful.

[egparedes]

LGTM. I have a couple of final questions, but they are not blockers, so I'm already approving the PR now.

[egparedes]

Just a question: since this is only used in the mypy plugin, would it be possible to move these definitions into the the plugin module?

[DropD]

I tried that first, but it seems the type has to be reachable from the context that is being checked. The only safe place for it is in the same module that defines Dimension.

[egparedes]

I guess this is the reason why we need xarray in the typing_exports group, but what is this actually testing?

[DropD]

I ran into a weird case where mypy crashed while looking at a module inside xarray (during icon4py pre-commit), because "Dims was not defined". This is a minimal reproducer for that crash, I left it as a regression test.