Fix a few typing issues

Author: inducer

sumpy/codegen.py

@@ -23,9 +23,10 @@
 THE SOFTWARE.
 """
 
-
+from abc import ABC
 import logging
 import re
+from typing import ParamSpec
 
 import numpy as np
 from constantdict import constantdict
@@ -43,6 +44,9 @@
 logger = logging.getLogger(__name__)
 
 
+P = ParamSpec("P")
+
+
 __doc__ = """
 
 Conversion of :mod:`sympy` expressions to :mod:`loopy`
@@ -237,7 +241,8 @@ def register_optimization_preambles(loopy_knl, device):
 
 # {{{ custom mapper base classes
 
-class CSECachingIdentityMapper(IdentityMapper, CSECachingMapperMixin):
+class CSECachingIdentityMapper(
+                    IdentityMapper[P], CSECachingMapperMixin[Expression, P], ABC):
     pass
 
 

sumpy/expansion/m2l.py

@@ -25,12 +25,13 @@
 
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, ClassVar
+from typing import Any, ClassVar, cast
 
 import numpy as np
 
 import loopy as lp
 import pymbolic
+import pymbolic.primitives as p
 
 import sumpy.symbolic as sym
 from sumpy.tools import add_to_sac, matvec_toeplitz_upper_triangular
@@ -1075,8 +1076,8 @@ def loopy_translation_classes_dependent_data(tgt_expansion, src_expansion,
     for i in range(len(insns)):
         insn = insns[i]
         if isinstance(insn, lp.Assignment) and \
-                insn.assignee.name.startswith(vec_name):
-            idx = int(insn.assignee.name[len(vec_name):])
+                cast(p.Variable, insn.assignee).name.startswith(vec_name):
+            idx = int(cast(p.Variable, insn.assignee).name[len(vec_name):])
             insns[i] = lp.Assignment(
                 assignee=data[idx],
                 expression=insn.expression,

sumpy/tools.py

@@ -972,7 +972,7 @@ def run_opencl_fft(
         input_vec: Any,
         inverse: bool = False,
         wait_for: list[pyopencl.Event] | None = None
-    ) -> tuple[pyopencl.Event, Any]:
+    ) -> tuple[pyopencl.Event | MarkerBasedProfilingEvent, Any]:
     """Runs an FFT on input_vec and returns a :class:`MarkerBasedProfilingEvent`
     that indicate the end and start of the operations carried out and the output
     vector.