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/action@v3.0.1
+
+  typing-exports:
+    needs: [get-python-versions]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ fromJSON(needs.get-python-versions.outputs.python-versions) }}
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        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,12 @@ typing = [
   'types-docutils>=0.21.0',
   'types-pytz>=2024.2.0'
 ]
+typing_exports = [
+  {include-group = "typing"},
+  'types-six',
+  'pytest-mypy-plugins',
+  "xarray"
+]
 
 # -- Standard project description options (PEP 621) --
 [project]
@@ -187,12 +193,13 @@ source_pkgs = ['gt4py']
 [tool.mypy]
 disallow_incomplete_defs = true
 disallow_untyped_defs = true
-exclude = ['^setup\.py$', 'build/.*$', 'ci/*.$', 'docs/.*$', 'tests/.*$']
+exclude = ['^setup\.py$', 'build/.*$', 'ci/*.$', 'docs/.*$', '^tests/.*$']
 ignore_missing_imports = true
 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 +268,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,56 @@
     "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",
 ]

src/gt4py/next/common.py

@@ -109,6 +109,24 @@ def __sub__(self, offset: int) -> Connectivity:
         return self + (-offset)
 
 
+if TYPE_CHECKING:
+
+    @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 +790,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 +1024,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 +1247,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(