Configure, pass mypy, mark typed

Author: inducer

.github/workflows/ci.yml

@@ -44,6 +44,23 @@ jobs:
 
                 run_pylint pymbolic test/test_*.py
 
+    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_venv
+                python -m pip install mypy numpy
+                ./run-mypy.sh
+
     pytest:
         name: Pytest on Py${{ matrix.python-version }}
         runs-on: ubuntu-latest

.gitlab-ci.yml

@@ -40,6 +40,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
+    ./run-mypy.sh
+  tags:
+  - python3
+  except:
+  - tags
+
 Documentation:
   script:
   - EXTRA_INSTALL="numpy sympy"

doc/conf.py

@@ -22,4 +22,6 @@
     "numpy": ("https://numpy.org/doc/stable/", None),
     "python": ("https://docs.python.org/3", None),
     "sympy": ("https://docs.sympy.org/dev/", None),
+    "typing_extensions":
+        ("https://typing-extensions.readthedocs.io/en/latest/", None),
 }

pymbolic/imperative/statement.py

@@ -24,6 +24,7 @@
 
 from sys import intern
 from pytools import RecordWithoutPickling
+from pymbolic.typing import not_none
 
 
 # {{{ statemetn classes
@@ -88,7 +89,7 @@ def get_dependency_mapper(self, include_calls="descend_args"):
 # {{{ statement with condition
 
 class ConditionalStatement(Statement):
-    __doc__ = Statement.__doc__ + """
+    __doc__ = not_none(Statement.__doc__) + """
     .. attribute:: condition
 
        The instruction condition as a :mod:`pymbolic` expression (`True` if the

pymbolic/interop/ast.py

@@ -288,7 +288,9 @@ def map_constant(self, expr: ScalarT) -> ast.expr:
     def map_call(self, expr: p.Call) -> ast.expr:
         return ast.Call(
             func=self.rec(expr.function),
-            args=[self.rec(param) for param in expr.parameters])
+            args=[self.rec(param) for param in expr.parameters],
+            keywords=[],
+        )
 
     def map_call_with_kwargs(self, expr) -> ast.expr:
         return ast.Call(
@@ -354,7 +356,7 @@ def map_bitwise_and(self, expr) -> ast.expr:
                                            ast.BitAnd())
 
     def map_logical_not(self, expr) -> ast.expr:
-        return ast.UnaryOp(self.rec(expr.child), ast.Not())
+        return ast.UnaryOp(ast.Not(), self.rec(expr.child))
 
     def map_logical_or(self, expr) -> ast.expr:
         return ast.BoolOp(ast.Or(), [self.rec(child)
@@ -376,6 +378,7 @@ def map_if(self, expr: p.If) -> ast.expr:
                          orelse=self.rec(expr.else_))
 
     def map_nan(self, expr: p.NaN) -> ast.expr:
+        assert expr.data_type is not None
         if isinstance(expr.data_type(float("nan")), float):
             return ast.Call(
                 ast.Name(id="float"),

pymbolic/interop/matchpy/__init__.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 """
 Interoperability with :mod:`matchpy.functions` for pattern-matching and
 term-rewriting.
@@ -38,22 +41,21 @@
 """
 
 
-import numpy as np
 import abc
 import pymbolic.primitives as p
 
-from typing import (Union, ClassVar, Optional, Iterator, Mapping,
-                    Generic, TypeVar, Tuple, Iterable, Callable)
+from typing_extensions import TypeAlias
+from typing import (ClassVar, Iterator, Mapping,
+                    Generic, TypeVar, Iterable, Callable)
 from dataclasses import dataclass, fields, field
 
 from matchpy import (Operation, Arity, Expression, Atom as BaseAtom,
                      Wildcard as BaseWildcard, ReplacementRule)
-from numbers import Number
 from functools import partial
+from pymbolic.typing import ScalarT
 
 
-ScalarT = Union[Number, int, np.bool_, bool]
-ExprT = Expression
+ExprT: TypeAlias = Expression
 ConstantT = TypeVar("ConstantT")
 ToMatchpyT = Callable[[p.Expression], ExprT]
 FromMatchpyT = Callable[[ExprT], p.Expression]
@@ -70,7 +72,7 @@
 @op_dataclass
 class _Constant(BaseAtom, Generic[ConstantT]):
     value: ConstantT
-    variable_name: Optional[str] = None
+    variable_name: str | None = None
 
     @property
     def head(self):
@@ -89,7 +91,7 @@ def __lt__(self, other):
 
 @op_dataclass
 class Scalar(_Constant[ScalarT]):
-    _mapper_method: [str] = "map_scalar"
+    _mapper_method: str = "map_scalar"
 
 
 @op_dataclass
@@ -104,8 +106,8 @@ class ComparisonOp(_Constant[str]):
 
 @op_dataclass
 class TupleOp(Operation):
-    _operands: Tuple[ExprT, ...]
-    variable_name: Optional[str] = non_operand_field(default=None)
+    _operands: tuple[ExprT, ...]
+    variable_name: str | None = non_operand_field(default=None)
 
     arity: ClassVar[Arity] = Arity.variadic
     name: ClassVar[str] = "tuple"
@@ -129,7 +131,7 @@ def variable_name(self):
         pass
 
     @property
-    def operands(self) -> Tuple[Expression]:
+    def operands(self) -> tuple[Expression, ...]:
         return tuple(getattr(self, field.name)
                      for field in fields(self)
                      if not field.metadata.get("not_an_operand", False))
@@ -150,15 +152,15 @@ def name(self) -> str:
 class Variable(PymbolicOp):
     id: Id
     arity: ClassVar[Arity] = Arity.unary
-    variable_name: Optional[str] = non_operand_field(default=None)
+    variable_name: str | None = non_operand_field(default=None)
     _mapper_method: ClassVar[str] = "map_variable"
 
 
 @op_dataclass
 class Call(PymbolicOp):
     function: ExprT
     args: TupleOp
-    variable_name: Optional[str] = non_operand_field(default=None)
+    variable_name: str | None = non_operand_field(default=None)
 
     arity: ClassVar[Arity] = Arity.binary
     _mapper_method: ClassVar[str] = "map_call"
@@ -168,7 +170,7 @@ class Call(PymbolicOp):
 class Subscript(PymbolicOp):
     aggregate: ExprT
     indices: TupleOp
-    variable_name: Optional[str] = non_operand_field(default=None)
+    variable_name: str | None = non_operand_field(default=None)
 
     arity: ClassVar[Arity] = Arity.binary
     _mapper_method: ClassVar[str] = "map_subscript"
@@ -182,7 +184,7 @@ class _BinaryOp(PymbolicOp):
     x2: ExprT
 
     arity: ClassVar[Arity] = Arity.binary
-    variable_name: Optional[str] = non_operand_field(default=None)
+    variable_name: str | None = non_operand_field(default=None)
 
 
 @op_dataclass
@@ -224,8 +226,8 @@ class RightShift(_BinaryOp):
 
 @variadic_op_dataclass
 class _VariadicCommAssocOp(PymbolicOp):
-    children: Tuple[ExprT, ...]
-    variable_name: Optional[str] = non_operand_field(default=None)
+    children: tuple[ExprT, ...]
+    variable_name: str | None = non_operand_field(default=None)
 
     commutative: ClassVar[bool] = True
     associative: ClassVar[bool] = True
@@ -236,7 +238,7 @@ def __init__(self, *children: ExprT, variable_name=None):
         object.__setattr__(self, "variable_name", variable_name)
 
     @property
-    def operands(self) -> Tuple[ExprT, ...]:
+    def operands(self) -> tuple[ExprT, ...]:
         return self.children
 
 
@@ -283,7 +285,7 @@ class BitwiseXor(_VariadicCommAssocOp):
 class _UnaryOp(PymbolicOp):
     x: ExprT
     arity: ClassVar[Arity] = Arity.unary
-    variable_name: Optional[str] = non_operand_field(default=None)
+    variable_name: str | None = non_operand_field(default=None)
 
 
 @op_dataclass
@@ -303,7 +305,7 @@ class Comparison(PymbolicOp):
     left: ExprT
     operator: ComparisonOp
     right: ExprT
-    variable_name: Optional[str] = non_operand_field(default=None)
+    variable_name: str | None = non_operand_field(default=None)
 
     arity: ClassVar[Arity] = Arity.ternary
     _mapper_method: ClassVar[str] = "map_comparison"
@@ -314,7 +316,7 @@ class If(PymbolicOp):
     condition: ExprT
     then: ExprT
     else_: ExprT
-    variable_name: Optional[str] = non_operand_field(default=None)
+    variable_name: str | None = non_operand_field(default=None)
 
     arity: ClassVar[Arity] = Arity.ternary
     _mapper_method: ClassVar[str] = "map_if"
@@ -325,16 +327,16 @@ class Wildcard(BaseWildcard):
     # {{{ FIXME: This should go into matchpy itself.
 
     @classmethod
-    def dot(cls, name=None) -> "Wildcard":
+    def dot(cls, name=None) -> Wildcard:
         return cls(min_count=1, fixed_size=True, variable_name=name)
 
     @classmethod
-    def star(cls, name=None) -> "Wildcard":
+    def star(cls, name=None) -> Wildcard:
         # FIXME: This should go into matchpy itself.
         return cls(min_count=0, fixed_size=False, variable_name=name)
 
     @classmethod
-    def plus(cls, name=None) -> "Wildcard":
+    def plus(cls, name=None) -> Wildcard:
         # FIXME: This should go into matchpy itself.
         return cls(min_count=1, fixed_size=False, variable_name=name)
 
@@ -343,7 +345,7 @@ def plus(cls, name=None) -> "Wildcard":
 # }}}
 
 
-def _get_operand_at_path(expr: PymbolicOp, path: Tuple[int, ...]) -> PymbolicOp:
+def _get_operand_at_path(expr: PymbolicOp, path: tuple[int, ...]) -> PymbolicOp:
     result = expr
 
     while path:
@@ -355,9 +357,9 @@ def _get_operand_at_path(expr: PymbolicOp, path: Tuple[int, ...]) -> PymbolicOp:
 
 def match(subject: p.Expression,
           pattern: p.Expression,
-          to_matchpy_expr: Optional[ToMatchpyT] = None,
-          from_matchpy_expr: Optional[FromMatchpyT] = None
-          ) -> Iterator[Mapping[str, Union[p.Expression, ScalarT]]]:
+          to_matchpy_expr: ToMatchpyT | None = None,
+          from_matchpy_expr: FromMatchpyT | None = None
+          ) -> Iterator[Mapping[str, p.Expression | ScalarT]]:
     from matchpy import match, Pattern
     from .tofrom import (ToMatchpyExpressionMapper,
                          FromMatchpyExpressionMapper)
@@ -378,10 +380,10 @@ def match(subject: p.Expression,
 
 def match_anywhere(subject: p.Expression,
                    pattern: p.Expression,
-                   to_matchpy_expr: Optional[ToMatchpyT] = None,
-                   from_matchpy_expr: Optional[FromMatchpyT] = None
-                   ) -> Iterator[Tuple[Mapping[str, Union[p.Expression, ScalarT]],
-                                       Union[p.Expression, ScalarT]]
+                   to_matchpy_expr: ToMatchpyT | None = None,
+                   from_matchpy_expr: FromMatchpyT | None = None
+                   ) -> Iterator[tuple[Mapping[str, p.Expression | ScalarT],
+                                       p.Expression | ScalarT]
                                  ]:
     from matchpy import match_anywhere, Pattern
     from .tofrom import (ToMatchpyExpressionMapper,
@@ -404,8 +406,8 @@ def match_anywhere(subject: p.Expression,
 
 def make_replacement_rule(pattern: p.Expression,
                           replacement: Callable[..., p.Expression],
-                          to_matchpy_expr: Optional[ToMatchpyT] = None,
-                          from_matchpy_expr: Optional[FromMatchpyT] = None
+                          to_matchpy_expr: ToMatchpyT | None = None,
+                          from_matchpy_expr: FromMatchpyT | None = None
                           ) -> ReplacementRule:
     """
     Returns a :class:`matchpy.functions.ReplacementRule` from the objects
@@ -429,9 +431,9 @@ def make_replacement_rule(pattern: p.Expression,
 
 def replace_all(expression: p.Expression,
                 rules: Iterable[ReplacementRule],
-                to_matchpy_expr: Optional[ToMatchpyT] = None,
-                from_matchpy_expr: Optional[FromMatchpyT] = None
-                ) -> Union[p.Expression, Tuple[p.Expression, ...]]:
+                to_matchpy_expr: ToMatchpyT | None = None,
+                from_matchpy_expr: FromMatchpyT | None = None
+                ) -> p.Expression | tuple[p.Expression, ...]:
     import collections.abc as abc
     from .tofrom import (ToMatchpyExpressionMapper,
                          FromMatchpyExpressionMapper)

pymbolic/interop/maxima.py

@@ -36,14 +36,14 @@
 
 import re
 from sys import intern
-from typing import ClassVar, List, Tuple
+from typing import ClassVar
 
 import numpy as np
 
 import pytools
 
 from pymbolic.mapper.stringifier import StringifyMapper
-from pymbolic.parser import Parser as ParserBase, FinalizedTuple
+from pymbolic.parser import Parser as ParserBase, FinalizedTuple, LexTable
 
 
 IN_PROMPT_RE = re.compile(br"\(%i([0-9]+)\) ")
@@ -95,7 +95,7 @@ class MaximaParser(ParserBase):
     imag_unit = intern("imag_unit")
     euler_number = intern("euler_number")
 
-    lex_table: ClassVar[List[Tuple[str, str]]] = [
+    lex_table: ClassVar[LexTable] = [
             (power_sym, pytools.lex.RE(r"\^")),
             (imag_unit, pytools.lex.RE(r"%i")),
             (euler_number, pytools.lex.RE(r"%e")),

pymbolic/mapper/constant_folder.py

@@ -95,7 +95,8 @@ class ConstantFoldingMapper(
             IdentityMapper.map_common_subexpression
 
 
-class CommutativeConstantFoldingMapper(
+# Yes, map_product incompatible: missing *args, **kwargs
+class CommutativeConstantFoldingMapper(    # type: ignore[misc]
         CSECachingMapperMixin,
         CommutativeConstantFoldingMapperBase,
         IdentityMapper):