Get mypy-clean, enable mypy in CI

Also, required from __future__ import annotations

.github/workflows/ci.yml

@@ -25,6 +25,23 @@ jobs:
         -   uses: actions/checkout@v4
         -   uses: crate-ci/typos@master
 
+    mypy:
+        name: Mypy
+        runs-on: ubuntu-latest
+        steps:
+        -   uses: actions/checkout@v4
+        -
+            uses: actions/setup-python@v5
+            with:
+                python-version: '3.x'
+        -   name: "Main Script"
+            run: |
+                curl -L -O https://tiker.net/ci-support-v0
+                . ./ci-support-v0
+                build_py_project_in_conda_env
+                python -m pip install mypy
+                mypy $(get_proj_name)
+
     pylint:
         name: Pylint
         runs-on: ubuntu-latest

.gitlab-ci.yml

@@ -140,6 +140,18 @@ Pylint:
   except:
   - tags
 
+Mypy:
+  script: |
+    curl -L -O https://tiker.net/ci-support-v0
+    . ./ci-support-v0
+    build_py_project_in_venv
+    python -m pip install mypy
+    mypy $(get_proj_name)
+  tags:
+  - python3
+  except:
+  - tags
+
 Downstream:
   parallel:
     matrix:

benchmarks/bench_translations.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 
 import numpy as np

pyproject.toml

@@ -16,6 +16,9 @@ extend-select = [
     "UP",  # pyupgrade
     "RUF", # ruff
     "W",   # pycodestyle
+
+    # TODO
+    # "SIM",
 ]
 extend-ignore = [
     "C90",  # McCabe complexity
@@ -42,6 +45,12 @@ known-local-folder = [
     "sumpy",
 ]
 lines-after-imports = 2
+required-imports = ["from __future__ import annotations"]
+
+[tool.ruff.lint.per-file-ignores]
+"doc/**/*.py" = ["I002"]
+"examples/**/*.py" = ["I002"]
+"setup.py" = ["I002"]
 
 [tool.typos.default]
 extend-ignore-re = [
@@ -61,3 +70,23 @@ extend-exclude = [
   "contrib/*/*.ipynb",
   "notes/*/*.eps",
 ]
+
+[tool.mypy]
+warn_unused_ignores = true
+
+[[tool.mypy.overrides]]
+module = [
+    "pyopencl.*",
+    "pymbolic.*",
+    "loopy.*",
+    "symengine.*",
+    "boxtree.*",
+    "pyvkfft.*",
+    "pyfmmlib.*",
+    "mayavi.*",
+    "scipy.*",
+    "sympy.*",
+    "matplotlib.*",
+    "pyvisfile.*",
+]
+ignore_missing_imports = true

sumpy/__init__.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = "Copyright (C) 2013 Andreas Kloeckner"
 
 __license__ = """
@@ -21,7 +24,9 @@
 """
 
 import os
+from typing import Hashable
 
+import loopy as lp
 from pytools.persistent_dict import WriteOncePersistentDict
 
 from sumpy.e2e import (
@@ -56,8 +61,8 @@
 ]
 
 
-code_cache = WriteOncePersistentDict("sumpy-code-cache-v6-"+VERSION_TEXT,
-                                     safe_sync=False)
+code_cache: WriteOncePersistentDict[Hashable, lp.TranslationUnit] = \
+        WriteOncePersistentDict("sumpy-code-cache-v6-"+VERSION_TEXT, safe_sync=False)
 
 
 # {{{ optimization control

sumpy/array_context.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = "Copyright (C) 2022 Alexandru Fikl"
 
 __license__ = """

sumpy/assignment_collection.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = "Copyright (C) 2012 Andreas Kloeckner"
 
 __license__ = """

sumpy/codegen.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = "Copyright (C) 2012 Andreas Kloeckner"
 
 __license__ = """

sumpy/cse.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = """
 Copyright (C) 2017 Matt Wala
 Copyright (C) 2006-2016 SymPy Development Team

sumpy/derivative_taker.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = """
 Copyright (C) 2012 Andreas Kloeckner
 Copyright (C) 2020 Isuru Fernando
@@ -338,7 +341,7 @@ def diff(self, mi, q=0):
 
 # {{{ DifferentiatedExprDerivativeTaker
 
-DerivativeCoeffDict = Dict[Tuple[int], Any]
+DerivativeCoeffDict = Dict[Tuple[int, ...], Any]
 
 
 @tag_dataclass
@@ -390,7 +393,7 @@ def diff_derivative_coeff_dict(derivative_coeff_dict: DerivativeCoeffDict,
     and return a new derivative transformation dictionary.
     """
     from collections import defaultdict
-    new_derivative_coeff_dict = defaultdict(lambda: 0)
+    new_derivative_coeff_dict: DerivativeCoeffDict = defaultdict(lambda: 0)
 
     for mi, coeff in derivative_coeff_dict.items():
         # In the case where we have x * u.diff(x), the result should

sumpy/distributed.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = "Copyright (C) 2022 Hao Gao"
 
 __license__ = """

sumpy/e2e.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = "Copyright (C) 2013 Andreas Kloeckner"
 
 __license__ = """

sumpy/e2p.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = "Copyright (C) 2013 Andreas Kloeckner"
 
 __license__ = """

sumpy/expansion/__init__.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = "Copyright (C) 2012 Andreas Kloeckner"
 
 __license__ = """
@@ -22,7 +25,7 @@
 
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, ClassVar, Dict, Hashable, List, Optional, Sequence, Tuple, Type
+from typing import Any, ClassVar, Hashable, Sequence
 
 import loopy as lp
 import pymbolic.primitives as prim
@@ -77,12 +80,12 @@ class ExpansionBase(ABC):
     .. automethod:: __eq__
     .. automethod:: __ne__
     """
-    init_arg_names = ("kernel", "order", "use_rscale")
+    init_arg_names: tuple[str, ...] = ("kernel", "order", "use_rscale")
 
     def __init__(self,
             kernel: Kernel,
             order: int,
-            use_rscale: Optional[bool] = None) -> None:
+            use_rscale: bool | None = None) -> None:
         # Don't be tempted to remove target derivatives here.
         # Line Taylor QBX can't do without them, because it can't
         # apply those derivatives to the expanded quantity.
@@ -125,7 +128,7 @@ def get_source_args(self):
     # {{{ abstract interface
 
     @abstractmethod
-    def get_coefficient_identifiers(self) -> List[Hashable]:
+    def get_coefficient_identifiers(self) -> list[Hashable]:
         """
         :returns: the identifiers of the coefficients that actually get stored.
         """
@@ -195,10 +198,10 @@ def loopy_evaluator(self, kernels: Sequence[Kernel]) -> lp.TranslationUnit:
 
     # {{{ copy
 
-    def with_kernel(self, kernel: Kernel) -> "ExpansionBase":
+    def with_kernel(self, kernel: Kernel) -> ExpansionBase:
         return type(self)(kernel, self.order, self.use_rscale)
 
-    def copy(self, **kwargs) -> "ExpansionBase":
+    def copy(self, **kwargs) -> ExpansionBase:
         new_kwargs = {
                 name: getattr(self, name)
                 for name in self.init_arg_names}
@@ -250,20 +253,20 @@ class ExpansionTermsWrangler(ABC):
 
     .. automethod:: get_full_coefficient_identifiers
     """
-    init_arg_names = ("order", "dim", "max_mi")
+    init_arg_names: tuple[str, ...] = ("order", "dim", "max_mi")
 
     def __init__(self,
             order: int,
             dim: int,
-            max_mi: Optional[int] = None) -> None:
+            max_mi: tuple[int, ...] | None = None) -> None:
         self.order = order
         self.dim = dim
         self.max_mi = max_mi
 
     # {{{ abstract interface
 
     @abstractmethod
-    def get_coefficient_identifiers(self) -> List[Hashable]:
+    def get_coefficient_identifiers(self) -> list[tuple[int, ...]]:
         pass
 
     @abstractmethod
@@ -279,7 +282,7 @@ def get_stored_mpole_coefficients_from_full(self,
     # }}}
 
     @memoize_method
-    def get_full_coefficient_identifiers(self) -> List[Hashable]:
+    def get_full_coefficient_identifiers(self) -> list[Hashable]:
         """
         Returns identifiers for every coefficient in the complete expansion.
         """
@@ -311,7 +314,7 @@ def copy(self, **kwargs):
 
     # {{{ hyperplane helpers
 
-    def _get_mi_hyperpplanes(self) -> List[Tuple[int, int]]:
+    def _get_mi_hyperpplanes(self) -> list[tuple[int, int]]:
         r"""
         Coefficient storage is organized into "hyperplanes" in multi-index
         space. Potentially only a subset of these hyperplanes contain
@@ -332,7 +335,7 @@ def _get_mi_hyperpplanes(self) -> List[Tuple[int, int]]:
 
     @memoize_method
     def _split_coeffs_into_hyperplanes(
-            self) -> List[Tuple[int, List[Tuple[int, ...]]]]:
+            self) -> list[tuple[int, list[tuple[int, ...]]]]:
         r"""
         This splits the coefficients into :math:`O(p)` disjoint sets
         so that for each set, all the identifiers have the form,
@@ -511,7 +514,7 @@ class LinearPDEBasedExpansionTermsWrangler(ExpansionTermsWrangler):
 
     def __init__(self,
             order: int, dim: int, knl: Kernel,
-            max_mi: Optional[int] = None) -> None:
+            max_mi: tuple[int, ...] | None = None) -> None:
         r"""
         :param order: order of the expansion
         :param dim: number of dimensions
@@ -592,7 +595,7 @@ def mi_key(ident):
 
         return mi_key, axis_permutation
 
-    def _get_mi_hyperpplanes(self) -> List[Tuple[int, int]]:
+    def _get_mi_hyperpplanes(self) -> list[tuple[int, int]]:
         mis = self.get_full_coefficient_identifiers()
         mi_to_index = {mi: i for i, mi in enumerate(mis)}
 
@@ -804,8 +807,8 @@ def get_projection_matrix(self, rscale):
 # FIXME: This is called an expansion but doesn't inherit from ExpansionBase?
 
 class VolumeTaylorExpansionMixin:
-    expansion_terms_wrangler_class: ClassVar[Type[ExpansionTermsWrangler]]
-    expansion_terms_wrangler_cache: ClassVar[Dict[Hashable, Any]] = {}
+    expansion_terms_wrangler_class: ClassVar[type[ExpansionTermsWrangler]]
+    expansion_terms_wrangler_cache: ClassVar[dict[Hashable, Any]] = {}
 
     @classmethod
     def get_or_make_expansion_terms_wrangler(cls, *key):