fix[next]: Reduce type ignores in client code

.github/workflows/code-quality.yml

@@ -40,3 +40,24 @@ jobs:
 
       - name: "Run pre-commit"
         uses: pre-commit/[email protected]
+
+  typing-exports:
+    needs: [get-python-versions]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Set up Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: ${{ fromJSON(needs.get-python-versions.outputs.python-versions) }}
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          cache-dependency-glob: "uv.lock"
+
+      - name: "Run typing checks on exported entities"
+        run: |
+          ./noxfile.py -s 'test_typing_exports-${{ fromJson(needs.get-python-versions.outputs.python-versions) }}'

noxfile.py

@@ -330,5 +330,21 @@ def test_storage(
     )
 
 
+@nox.session(python=PYTHON_VERSIONS, tags=["next"])
+def test_typing_exports(session: nox.Session) -> None:
+    """Test GT4Py usability in a typed client context."""
+    install_session_venv(session, extras=["standard"], groups=["test", "typing_exports"])
+
+    session.run(
+        "pytest",
+        "-sv",
+        "--mypy-testing-base",
+        "typing_tests",
+        "--mypy-only-local-stub",
+        "typing_tests",
+        *session.posargs,
+    )
+
+
 if __name__ == "__main__":
     nox.main()

pyproject.toml

@@ -64,6 +64,13 @@ typing = [
   'types-docutils>=0.21.0',
   'types-pytz>=2024.2.0'
 ]
+typing_exports = [
+  # to test typing with gt4py in downstream code
+  {include-group = "typing"},
+  'types-six',  # can not let mypy auto-install types as that leads to unexpected stderr output (which means test failure)
+  'pytest-mypy-plugins',  # pytest plugin for running mypy on code snippets
+  "xarray"  # one of the regression tests requires xarray
+]
 
 # -- Standard project description options (PEP 621) --
 [project]
@@ -193,6 +200,7 @@ implicit_optional = false
 implicit_reexport = false
 install_types = true
 namespace_packages = false
+plugins = ['gt4py.next.type_system.mypy_plugin']
 # pretty = true
 show_column_numbers = true
 show_error_codes = true
@@ -261,6 +269,12 @@ implicit_reexport = true
 # factory-boy is broken, see https://github.com/FactoryBoy/factory_boy/pull/1114
 module = "factory.*"
 
+[[tool.mypy.overrides]]
+disallow_incomplete_defs = false
+disallow_untyped_defs = false
+ignore_errors = false
+module = "typing_tests.test_next_exports"
+
 # -- pytest --
 [tool.pytest]
 

src/gt4py/_core/definitions.py

@@ -166,7 +166,7 @@ def dtype_kind(
 
 
 @overload
-def dtype_kind(sc_type: Type[UnsignedIntT]) -> Literal[DTypeKind.UINT]: ...
+def dtype_kind(sc_type: Type[UnsignedIntT]) -> Literal[DTypeKind.UINT]: ...  # type: ignore[overload-cannot-match] # precision blurring from mypy plugin seems to interfere
 
 
 @overload

src/gt4py/next/__init__.py

@@ -39,8 +39,59 @@
 )
 from .ffront import fbuiltins
 from .ffront.decorator import field_operator, program, scan_operator
-from .ffront.fbuiltins import *  # noqa: F403 [undefined-local-with-import-star]  explicitly reexport all from fbuiltins.__all__
-from .ffront.fbuiltins import FieldOffset
+from .ffront.fbuiltins import (
+    FieldOffset,
+    IndexType,
+    abs,  # noqa: A004 # shadowing
+    arccos,
+    arccosh,
+    arcsin,
+    arcsinh,
+    arctan,
+    arctanh,
+    astype,
+    bool,  # noqa: A004 # shadowing
+    broadcast,
+    cbrt,
+    ceil,
+    cos,
+    cosh,
+    exp,
+    float,  # noqa: A004 # shadowing
+    float32,
+    float64,
+    floor,
+    fmod,
+    gamma,
+    int,  # noqa: A004 # shadowing
+    int8,
+    int16,
+    int32,
+    int64,
+    isfinite,
+    isinf,
+    isnan,
+    log,
+    max_over,
+    maximum,
+    min_over,
+    minimum,
+    neg,
+    neighbor_sum,
+    power,
+    sin,
+    sinh,
+    sqrt,
+    tan,
+    tanh,
+    trunc,
+    tuple,  # noqa: A004 # shadowing
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+    where,
+)
 from .otf.compiled_program import wait_for_compilation
 from .program_processors.runners.gtfn import (
     run_gtfn_cached as gtfn_cpu,
@@ -49,7 +100,7 @@
 from .program_processors.runners.roundtrip import default as itir_python
 
 
-__all__ = [
+__all__ = [  # noqa: RUF022 [unsorted-dunder-all]
     # submodules
     "common",
     "ffront",
@@ -91,5 +142,60 @@
     "gtfn_cpu",
     "gtfn_gpu",
     "itir_python",
-    *fbuiltins.__all__,
+    # from fbuiltins
+    "FieldOffset",
+    "IndexType",
+    "abs",
+    "arccos",
+    "arccosh",
+    "arcsin",
+    "arcsinh",
+    "arctan",
+    "arctanh",
+    "astype",
+    "bool",
+    "broadcast",
+    "cbrt",
+    "ceil",
+    "cos",
+    "cosh",
+    "exp",
+    "float",
+    "float32",
+    "float64",
+    "floor",
+    "fmod",
+    "gamma",
+    "int",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "isfinite",
+    "isinf",
+    "isnan",
+    "log",
+    "max_over",
+    "min_over",
+    "maximum",
+    "minimum",
+    "neg",
+    "neighbor_sum",
+    "power",
+    "sin",
+    "sinh",
+    "sqrt",
+    "tan",
+    "tanh",
+    "trunc",
+    "tuple",
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+    "where",
 ]
+
+assert not (diff := set(fbuiltins.__all__) - set(__all__)), (
+    f"public fbuiltins member(s) not exported: {diff}"
+)

src/gt4py/next/common.py

@@ -109,6 +109,28 @@ def __sub__(self, offset: int) -> Connectivity:
         return self + (-offset)
 
 
+if TYPE_CHECKING:
+    # These exist as on-the fly replacements for Dimension instances
+    # (which are not types) during typechecking (with mypy). We can
+    # track up to four distinct dimensions at a time, everything beyond
+    # becomes AnyDim
+
+    @dataclasses.dataclass(frozen=True)
+    class _DimA(Dimension): ...
+
+    @dataclasses.dataclass(frozen=True)
+    class _DimB(Dimension): ...
+
+    @dataclasses.dataclass(frozen=True)
+    class _DimC(Dimension): ...
+
+    @dataclasses.dataclass(frozen=True)
+    class _DimD(Dimension): ...
+
+    @dataclasses.dataclass(frozen=True)
+    class _AnyDim(Dimension): ...
+
+
 class Infinity(enum.Enum):
     """Describes an unbounded `UnitRange`."""
 
@@ -772,6 +794,18 @@ def __rtruediv__(self, other: Field | core_defs.ScalarT) -> Field: ...
     @abc.abstractmethod
     def __pow__(self, other: Field | core_defs.ScalarT) -> Field: ...
 
+    @abc.abstractmethod
+    def __lt__(self, other: Field | core_defs.ScalarT) -> Field[Any, bool]: ...
+
+    @abc.abstractmethod
+    def __le__(self, other: Field | core_defs.ScalarT) -> Field[Any, bool]: ...
+
+    @abc.abstractmethod
+    def __gt__(self, other: Field | core_defs.ScalarT) -> Field[Any, bool]: ...
+
+    @abc.abstractmethod
+    def __ge__(self, other: Field | core_defs.ScalarT) -> Field[Any, bool]: ...
+
     @abc.abstractmethod
     def __and__(self, other: Field | core_defs.ScalarT) -> Field:
         """Only defined for `Field` of value type `bool`."""
@@ -994,36 +1028,48 @@ def __ne__(self, other: Any) -> Never:
     def __add__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
-    def __radd__(self, other: Field | core_defs.IntegralScalar) -> Never:  # type: ignore[misc] # Forward operator not callalbe
+    def __radd__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
     def __sub__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
-    def __rsub__(self, other: Field | core_defs.IntegralScalar) -> Never:  # type: ignore[misc] # Forward operator not callalbe
+    def __rsub__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
     def __mul__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
-    def __rmul__(self, other: Field | core_defs.IntegralScalar) -> Never:  # type: ignore[misc] # Forward operator not callalbe
+    def __rmul__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
     def __truediv__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
-    def __rtruediv__(self, other: Field | core_defs.IntegralScalar) -> Never:  # type: ignore[misc] # Forward operator not callalbe
+    def __rtruediv__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
     def __floordiv__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
-    def __rfloordiv__(self, other: Field | core_defs.IntegralScalar) -> Never:  # type: ignore[misc] # Forward operator not callalbe
+    def __rfloordiv__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
     def __pow__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
+    def __lt__(self, other: Field | core_defs.IntegralScalar) -> Never:
+        raise TypeError("'Connectivity' does not support this operation.")
+
+    def __le__(self, other: Field | core_defs.IntegralScalar) -> Never:
+        raise TypeError("'Connectivity' does not support this operation.")
+
+    def __gt__(self, other: Field | core_defs.IntegralScalar) -> Never:
+        raise TypeError("'Connectivity' does not support this operation.")
+
+    def __ge__(self, other: Field | core_defs.IntegralScalar) -> Never:
+        raise TypeError("'Connectivity' does not support this operation.")
+
     def __and__(self, other: Field | core_defs.IntegralScalar) -> Never:
         raise TypeError("'Connectivity' does not support this operation.")
 
@@ -1205,7 +1251,7 @@ def __gt_origin__(self) -> Never:
 
     @property
     def dtype(self) -> core_defs.DType[core_defs.IntegralScalar]:
-        return core_defs.Int32DType()  # type: ignore[return-value]
+        return core_defs.Int32DType()
 
     # This is a workaround to make this class concrete, since `codomain` is an
     # abstract property of the `Connectivity` Protocol.

src/gt4py/next/embedded/nd_array_field.py

@@ -797,11 +797,11 @@ def _hyperslice(
 # -- Specialized implementations for builtin operations on array fields --
 
 NdArrayField.register_builtin_func(
-    fbuiltins.abs,  # type: ignore[attr-defined]
+    fbuiltins.abs,
     NdArrayField.__abs__,
 )
 NdArrayField.register_builtin_func(
-    fbuiltins.power,  # type: ignore[attr-defined]
+    fbuiltins.power,
     NdArrayField.__pow__,
 )
 # TODO gamma
@@ -816,15 +816,15 @@ def _hyperslice(
     NdArrayField.register_builtin_func(getattr(fbuiltins, name), _make_builtin(name, name))
 
 NdArrayField.register_builtin_func(
-    fbuiltins.minimum,  # type: ignore[attr-defined]
+    fbuiltins.minimum,
     _make_builtin("minimum", "minimum"),
 )
 NdArrayField.register_builtin_func(
-    fbuiltins.maximum,  # type: ignore[attr-defined]
+    fbuiltins.maximum,
     _make_builtin("maximum", "maximum"),
 )
 NdArrayField.register_builtin_func(
-    fbuiltins.fmod,  # type: ignore[attr-defined]
+    fbuiltins.fmod,
     _make_builtin("fmod", "fmod"),
 )
 NdArrayField.register_builtin_func(fbuiltins.where, _make_builtin("where", "where"))
@@ -1159,7 +1159,7 @@ def _astype(field: common.Field | core_defs.ScalarT | tuple, type_: type) -> NdA
     raise AssertionError("This is the NdArrayField implementation of 'fbuiltins.astype'.")
 
 
-NdArrayField.register_builtin_func(fbuiltins.astype, _astype)
+NdArrayField.register_builtin_func(fbuiltins.astype, _astype)  # type: ignore[arg-type]  # because fbuiltins.astype is overloaded
 
 
 def _get_slices_from_domain_slice(