Improve kernel type annotations

doc/conf.py

@@ -2,8 +2,7 @@
 from urllib.request import urlopen
 
 
-_conf_url = \
-        "https://raw.githubusercontent.com/inducer/sphinxconfig/main/sphinxconfig.py"
+_conf_url = "https://raw.githubusercontent.com/inducer/sphinxconfig/main/sphinxconfig.py"
 with urlopen(_conf_url) as _inf:
     exec(compile(_inf.read(), _conf_url, "exec"), globals())
 
@@ -28,3 +27,17 @@
 nitpick_ignore_regex = [
     ["py:class", r"symengine\.(.+)"],  # :cry:
 ]
+
+sphinxconfig_missing_reference_aliases = {
+    # sympy
+    "sp.Matrix": "class:sympy.matrices.dense.DenseMatrix",
+    # pymbolic
+    "ArithmeticExpression": "obj:pymbolic.ArithmeticExpression",
+    "Expression": "obj:pymbolic.typing.Expression",
+    # sumpy
+    "ArithmeticExpr": "obj:sumpy.kernel.ArithmeticExpr",
+}
+
+
+def setup(app):
+    app.connect("missing-reference", process_autodoc_missing_reference)  # noqa: F821

sumpy/codegen.py

@@ -549,7 +549,12 @@ def map_constant(self, expr, rec_self=None):
 class VectorComponentRewriter(CSECachingIdentityMapper, CallExternalRecMapper):
     """For names in name_whitelist, turn ``a3`` into ``a[3]``."""
 
-    def __init__(self, name_whitelist=frozenset()):
+    name_whitelist: frozenset[str]
+
+    def __init__(self, name_whitelist: frozenset[str] | None = None) -> None:
+        if name_whitelist is None:
+            name_whitelist = frozenset()
+
         self.name_whitelist = name_whitelist
 
     def map_variable(self, expr, *args):

sumpy/derivative_taker.py

@@ -39,7 +39,7 @@
 """
 
 import logging
-from typing import Any
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -49,6 +49,9 @@
 from sumpy.tools import add_mi, add_to_sac
 
 
+if TYPE_CHECKING:
+    import sympy as sp
+
 logger = logging.getLogger(__name__)
 
 
@@ -341,7 +344,7 @@ def diff(self, mi, q=0):
 
 # {{{ DifferentiatedExprDerivativeTaker
 
-DerivativeCoeffDict = dict[tuple[int, ...], Any]
+DerivativeCoeffDict = dict[tuple[int, ...], int | float | complex | sym.Expr]
 
 
 @tag_dataclass
@@ -387,8 +390,10 @@ def diff(self, mi, save_intermediate=lambda x: x):
 
 # {{{ Helper functions
 
-def diff_derivative_coeff_dict(derivative_coeff_dict: DerivativeCoeffDict,
-        variable_idx, variables):
+def diff_derivative_coeff_dict(
+        derivative_coeff_dict: DerivativeCoeffDict,
+        variable_idx: int,
+        variables: sp.Matrix) -> DerivativeCoeffDict:
     """Differentiate a derivative transformation dictionary given by
     *derivative_coeff_dict* using the variable given by **variable_idx**
     and return a new derivative transformation dictionary.
@@ -397,15 +402,17 @@ def diff_derivative_coeff_dict(derivative_coeff_dict: DerivativeCoeffDict,
     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
-        # be x.diff(x) + x * u.diff(x, x)
+        # In the case where we have x * u.diff(x), the result should be
+        #       x.diff(x) + x * u.diff(x, x)
         # Calculate the first term by differentiating the coefficients
         new_coeff = sym.sympify(coeff).diff(variables[variable_idx])
         new_derivative_coeff_dict[mi] += new_coeff
+
         # Next calculate the second term by differentiating the derivatives
         new_mi = list(mi)
         new_mi[variable_idx] += 1
         new_derivative_coeff_dict[tuple(new_mi)] += coeff
+
     return {derivative: coeff for derivative, coeff in
             new_derivative_coeff_dict.items() if coeff != 0}
 

sumpy/e2e.py

@@ -387,7 +387,7 @@ def get_inner_loopy_kernel(self, result_dtype):
             ncoeff_src = len(self.src_expansion)
             ncoeff_tgt = len(self.tgt_expansion)
 
-        domains = []
+        domains: list[str] = []
         insns = self.get_translation_loopy_insns(result_dtype)
         tgt_coeffs = pymbolic.var("tgt_coeffs")
         for i in range(ncoeff_tgt):

sumpy/expansion/diff_op.py

@@ -113,7 +113,7 @@ def order(self) -> int:
             deg = max(deg, max(sum(ident.mi) for ident in eq))
         return deg
 
-    def __mul__(self, param: Number_ish) -> LinearPDESystemOperator:
+    def __mul__(self, param: Number_ish | sym.Expr) -> LinearPDESystemOperator:
         eqs: list[Mapping[DerivativeIdentifier, sp.Expr]] = []
         for eq in self.eqs:
             deriv_ident_to_coeff = {}