Field spec

.github/workflows/ci.yml

@@ -73,6 +73,24 @@ jobs:
                 python -m pip install mypy numpy
                 ./run-mypy.sh
 
+    basedpyright:
+        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_venv
+                pip install -e .[test]
+                python -m pip install numpy pexpect sympy scipy
+                python -m pip install basedpyright
+                basedpyright
+
     pytest:
         name: Pytest on Py${{ matrix.python-version }}
         runs-on: ubuntu-latest

pymbolic/algorithm.py

@@ -8,6 +8,23 @@
 .. autofunction:: sym_fft
 .. autofunction:: reduced_row_echelon_form
 .. autofunction:: solve_affine_equations_for
+
+References
+----------
+
+.. class:: NDArray
+
+    See :data:`numpy.typing.NDArray`.
+
+.. currentmodule:: np
+
+.. class:: generic
+
+    See :class:`numpy.generic`.
+
+.. class:: inexact
+
+    See :class:`numpy.inexact`.
 """
 
 from __future__ import annotations
@@ -45,6 +62,7 @@
 
 if TYPE_CHECKING:
     import numpy as np
+    from numpy.typing import NDArray
 
 
 if getattr(sys, "_BUILDING_SPHINX_DOCS", None):
@@ -312,26 +330,26 @@ def csr_matrix_multiply(S, x):  # noqa
 
 @overload
 def reduced_row_echelon_form(
-            mat: np.ndarray,
+            mat: NDArray[np.inexact],
             *, integral: bool | None = None,
-        ) -> np.ndarray:
+        ) -> NDArray[np.inexact]:
     ...
 
 
 @overload
 def reduced_row_echelon_form(
-            mat: np.ndarray,
-            rhs: np.ndarray,
+            mat: NDArray[np.inexact],
+            rhs: NDArray[np.inexact],
             *, integral: bool | None = None,
-        ) -> tuple[np.ndarray, np.ndarray]:
+        ) -> tuple[NDArray[np.inexact], NDArray[np.inexact]]:
     ...
 
 
 def reduced_row_echelon_form(
-            mat: np.ndarray,
-            rhs: np.ndarray | None = None,
+            mat: NDArray[np.inexact],
+            rhs: NDArray[np.inexact] | None = None,
             integral: bool | None = None,
-        ) -> tuple[np.ndarray, np.ndarray] | np.ndarray:
+        ) -> tuple[NDArray[np.inexact], NDArray[np.inexact]] | NDArray[np.inexact]:
     m, n = mat.shape
 
     mat = mat.copy()

pymbolic/geometric_algebra/__init__.py

@@ -27,7 +27,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable, Mapping, Sequence
 from dataclasses import dataclass
-from typing import Any, Generic, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast
 
 import numpy as np
 
@@ -37,6 +37,10 @@
 from pymbolic.typing import ArithmeticExpression, T
 
 
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+
 __doc__ = """
 See `Wikipedia <https://en.wikipedia.org/wiki/Geometric_algebra>`__ for an idea
 of what this is.
@@ -124,7 +128,6 @@
 .. literalinclude:: ../test/test_pymbolic.py
    :start-after: START_GA_TEST
    :end-before: END_GA_TEST
-
 """
 
 
@@ -195,15 +198,15 @@ class Space:
     basis_names: Sequence[str]
     "A sequence of names of basis vectors."
 
-    metric_matrix: np.ndarray
+    metric_matrix: NDArray[np.generic]
     """
     A *(dims, dims)*-shaped matrix, whose *(i, j)*-th entry represents the
     inner product of basis vector *i* and basis vector *j*.
     """
 
     def __init__(self,
                  basis: Sequence[str] | int | None = None,
-                 metric_matrix: np.ndarray | None = None) -> None:
+                 metric_matrix: NDArray[np.generic] | None = None) -> None:
         """
         :arg basis: A sequence of names of basis vectors, or an integer (the
             number of dimensions) to use the default names ``e0`` through ``eN``.
@@ -539,7 +542,7 @@ def __init__(
                 self,
                 data: (Mapping[int, CoeffT]
                        | Mapping[tuple[int, ...], CoeffT]
-                       | np.ndarray
+                       | NDArray[np.generic]
                        | CoeffT),
                 space: Space | None = None
             ) -> None:

pymbolic/mapper/__init__.py

@@ -44,6 +44,7 @@
 
 if TYPE_CHECKING:
     import numpy as np
+    from numpy.typing import NDArray
 
     from pymbolic.geometric_algebra import MultiVector
     from pymbolic.rational import Rational
@@ -121,16 +122,16 @@
 if TYPE_CHECKING:
     import numpy as np
 
-    def is_numpy_array(val) -> TypeIs[np.ndarray]:
+    def is_numpy_array(val: object) -> TypeIs[NDArray[np.generic]]:
         return isinstance(val, np.ndarray)
 else:
     try:
         import numpy as np
 
-        def is_numpy_array(val):
+        def is_numpy_array(val: object) -> bool:
             return isinstance(val, np.ndarray)
     except ImportError:
-        def is_numpy_array(ary):
+        def is_numpy_array(ary) -> bool:
             return False
 
 
@@ -200,7 +201,7 @@
             else:
                 return self.handle_unsupported_expression(expr, *args, **kwargs)
         else:
             return self.map_foreign(expr, *args, **kwargs)
 
     rec = __call__
 
@@ -298,7 +299,7 @@
         raise NotImplementedError
 
     def map_numpy_array(self,
-            expr: np.ndarray, *args: P.args, **kwargs: P.kwargs) -> ResultT:
+            expr: NDArray[np.generic], *args: P.args, **kwargs: P.kwargs) -> ResultT:
         raise NotImplementedError
 
     def map_left_shift(self,
@@ -602,9 +603,10 @@
         return self.combine(self.rec(child, *args, **kwargs) for child in expr)
 
     def map_numpy_array(self,
-                expr: np.ndarray, *args: P.args, **kwargs: P.kwargs
+                expr: NDArray[np.generic], *args: P.args, **kwargs: P.kwargs
             ) -> ResultT:
-        return self.combine(self.rec(el, *args, **kwargs) for el in expr.flat)
+        # FIXME Can the type-ignore be avoided?
+        return self.combine(self.rec(el, *args, **kwargs) for el in expr.flat)  # type: ignore[arg-type]
 
     def map_multivector(self,
                 expr: MultiVector[ArithmeticExpression],
@@ -680,7 +682,7 @@
         return set()
 
 
-class CachedCollector(CachedMapper, Collector):
+class CachedCollector(CachedMapper[Set[CollectedT], P], Collector[CollectedT, P]):
     pass
 
 # }}}
@@ -957,7 +959,7 @@
         return tuple(children)
 
     def map_numpy_array(self,
-                expr: np.ndarray, *args: P.args, **kwargs: P.kwargs
+                expr: NDArray[np.generic], *args: P.args, **kwargs: P.kwargs
             ) -> Expression:
 
         import numpy
@@ -1231,7 +1233,7 @@
         self.post_visit(expr, *args, **kwargs)
 
     def map_numpy_array(self,
-            expr: np.ndarray, *args: P.args, **kwargs: P.kwargs) -> None:
+            expr: NDArray[np.generic], *args: P.args, **kwargs: P.kwargs) -> None:
         if not self.visit(expr, *args, **kwargs):
             return
 
@@ -1445,7 +1447,7 @@
 # {{{ callback mapper
 
 # FIXME: Is it worth typing this?
-class CallbackMapper(Mapper):
+class CallbackMapper(Mapper[ResultT, P]):
     def __init__(self, function, fallback_mapper):
         self.function = function
         self.fallback_mapper = fallback_mapper

pymbolic/mapper/evaluator.py

@@ -44,6 +44,7 @@
     from collections.abc import Mapping
 
     import numpy as np
+    from numpy.typing import NDArray
 
     import pymbolic.primitives as p
     from pymbolic.geometric_algebra import MultiVector
@@ -54,7 +55,7 @@ class UnknownVariableError(Exception):
     pass
 
 
-class EvaluationMapper(Mapper[ResultT, []], CSECachingMapperMixin):
+class EvaluationMapper(Mapper[ResultT, []], CSECachingMapperMixin[ResultT, []]):
     """Example usage:
 
     .. doctest::
@@ -159,7 +160,7 @@ def map_logical_and(self, expr: p.LogicalAnd) -> bool:  # type: ignore[override]
     def map_list(self, expr: list[Expression]) -> ResultT:
         return [self.rec(child) for child in expr]  # type: ignore[return-value]
 
-    def map_numpy_array(self, expr: np.ndarray) -> ResultT:
+    def map_numpy_array(self, expr: NDArray[np.generic]) -> ResultT:
         import numpy
         result = numpy.empty(expr.shape, dtype=object)
         for i in numpy.ndindex(expr.shape):

pymbolic/mapper/stringifier.py

@@ -35,6 +35,7 @@
     from collections.abc import Sequence
 
     import numpy as np
+    from numpy.typing import NDArray
 
     from pymbolic.geometric_algebra import MultiVector
     from pymbolic.typing import Expression
@@ -594,7 +595,7 @@ def map_tuple(
 
     def map_numpy_array(
         self,
-        expr: np.ndarray,
+        expr: NDArray[np.generic],
         enclosing_prec: int,
         *args: P.args,
         **kwargs: P.kwargs,

pymbolic/primitives.py

@@ -23,6 +23,7 @@
 THE SOFTWARE.
 """
 
+import dataclasses
 import re
 from dataclasses import dataclass, fields
 from functools import partial
@@ -1077,7 +1078,7 @@ def {cls.__name__}_setstate(self, state):
 _T = TypeVar("_T")
 
 
-@dataclass_transform(frozen_default=True)
+@dataclass_transform(frozen_default=True, field_specifiers=(dataclasses.field,))
 def expr_dataclass(
             init: bool = True,
             eq: bool = True,

pyproject.toml

@@ -116,3 +116,17 @@ module = [
     "pexpect.*",
 ]
 ignore_missing_imports = true
+
+[tool.basedpyright]
+reportImplicitStringConcatenation = "none"
+reportUnnecessaryIsInstance = "none"
+reportUnusedCallResult = "none"
+reportExplicitAny = "none"
+
+# This reports even cycles that are qualified by 'if TYPE_CHECKING'. Not what
+# we care about at this moment.
+# https://github.com/microsoft/pyright/issues/746
+reportImportCycles = "none"
+pythonVersion = "3.10"
+pythonPlatform = "All"
+