Add type annotations around sumpy.fmm

doc/conf.py

@@ -13,15 +13,16 @@
 intersphinx_mapping = {
     "arraycontext": ("https://documen.tician.de/arraycontext/", None),
     "boxtree": ("https://documen.tician.de/boxtree/", None),
-    "namedisl": ("https://documen.tician.de/namedisl", None),
     "loopy": ("https://documen.tician.de/loopy/", None),
     "matplotlib": ("https://matplotlib.org/stable/", None),
+    "namedisl": ("https://documen.tician.de/namedisl", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
     "pymbolic": ("https://documen.tician.de/pymbolic/", None),
     "pyopencl": ("https://documen.tician.de/pyopencl/", None),
     "pytential": ("https://documen.tician.de/pytential/", None),
     "python": ("https://docs.python.org/3/", None),
     "pytools": ("https://documen.tician.de/pytools/", None),
+    "pyvkfft": ("https://pyvkfft.readthedocs.io/en/latest/", None),
     "sympy": ("https://docs.sympy.org/latest/", None),
 }
 
@@ -57,6 +58,11 @@
     # pymbolic
     "ArithmeticExpression": "obj:pymbolic.ArithmeticExpression",
     "Expression": "obj:pymbolic.typing.Expression",
+    # pyopencl
+    "CLArray": "class:pyopencl.array.Array",
+    "WaitList": "obj:pyopencl.WaitList",
+    # pyvkfft
+    "VkFFTApp": "class:pyvkfft.base.VkFFTApp",
     # namedisl
     "nisl.Set": "class:namedisl.Set",
     # loopy
@@ -71,6 +77,7 @@
     # sumpy
     "ArithmeticExpr": "obj:sumpy.kernel.ArithmeticExpr",
     "OptimizationPair": "obj:sumpy.cse.OptimizationPair",
+    "TranslationClassesInfo": "class:boxtree.translation_classes.TranslationClassesInfo",  # ruff: ignore[line-too-long]
 }
 
 

pyproject.toml

@@ -118,6 +118,7 @@ known-first-party = [
     "loopy",
     "pymbolic",
     "pyopencl",
+    "boxtree",
     "pytools",
 ]
 known-local-folder = [

sumpy/__init__.py

@@ -44,6 +44,7 @@
 
 
 if TYPE_CHECKING:
+    import types
     from collections.abc import Hashable
 
     import loopy as lp
@@ -72,15 +73,13 @@
 
 # {{{ optimization control
 
-OPT_ENABLED = True
-
 OPT_ENABLED = "SUMPY_NO_OPT" not in os.environ
 
 
-def set_optimization_enabled(flag):
+def set_optimization_enabled(flag: bool) -> None:
     """Set whether the :mod:`loopy` kernels should be optimized."""
     global OPT_ENABLED
-    OPT_ENABLED = flag
+    OPT_ENABLED = flag  # pyright: ignore[reportConstantRedefinition]
 
 # }}}
 
@@ -91,36 +90,55 @@ def set_optimization_enabled(flag):
     "SUMPY_NO_CACHE" not in os.environ
     and "CG_NO_CACHE" not in os.environ)
 
-NO_CACHE_KERNELS = tuple(os.environ.get("SUMPY_NO_CACHE_KERNELS",
-                                        "").split(","))
+NO_CACHE_KERNELS = tuple(os.environ.get("SUMPY_NO_CACHE_KERNELS", "").split(","))
 
 
-def set_caching_enabled(flag, no_cache_kernels=()):
+def set_caching_enabled(flag: bool, no_cache_kernels: tuple[str, ...] = ()) -> None:
     """Set whether :mod:`loopy` is allowed to use disk caching for its various
     code generation stages.
     """
     global CACHING_ENABLED, NO_CACHE_KERNELS
-    NO_CACHE_KERNELS = no_cache_kernels
-    CACHING_ENABLED = flag
+    NO_CACHE_KERNELS = no_cache_kernels  # pyright: ignore[reportConstantRedefinition]
+    CACHING_ENABLED = flag  # pyright: ignore[reportConstantRedefinition]
 
 
 class CacheMode:
     """A context manager for setting whether :mod:`sumpy` is allowed to use
     disk caches.
     """
 
-    def __init__(self, new_flag, new_no_cache_kernels=()):
+    new_flag: bool
+    previous_flag: bool | None
+
+    new_no_cache_kernels: tuple[str, ...]
+    previous_no_cache_kernels: tuple[str, ...] | None
+
+    def __init__(self,
+                 new_flag: bool,
+                 new_no_cache_kernels: tuple[str, ...] = ()) -> None:
         self.new_flag = new_flag
+        self.previous_flag = None
         self.new_no_cache_kernels = new_no_cache_kernels
+        self.previous_no_cache_kernels = None
+
+    def __enter__(self) -> None:
+        if self.previous_flag is not None or self.previous_no_cache_kernels is not None:
+            raise RuntimeError("cannot reuse the 'CacheMode' context manager")
 
-    def __enter__(self):
         self.previous_flag = CACHING_ENABLED
-        self.previous_kernels = NO_CACHE_KERNELS
+        self.previous_no_cache_kernels = NO_CACHE_KERNELS
         set_caching_enabled(self.new_flag, self.new_no_cache_kernels)
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        set_caching_enabled(self.previous_flag, self.previous_kernels)
-        del self.previous_flag
-        del self.previous_kernels
+    def __exit__(self,
+                 exc_type: type[BaseException] | None,
+                 exc_val: BaseException | None,
+                 exc_tb: types.TracebackType | None) -> None:
+        if self.previous_flag is None or self.previous_no_cache_kernels is None:
+            raise RuntimeError("cannot reuse the 'CacheMode' context manager")
+
+        set_caching_enabled(self.previous_flag, self.previous_no_cache_kernels)
+
+        self.previous_flag = None
+        self.previous_no_cache_kernels = None
 
 # }}}

sumpy/array_context.py

@@ -25,14 +25,14 @@
 
 from typing import TYPE_CHECKING, Any
 
-from boxtree.array_context import PyOpenCLArrayContext as PyOpenCLArrayContextBase
 from typing_extensions import override
 
 import loopy as lp
 from arraycontext.pytest import (
     _PytestPyOpenCLArrayContextFactoryWithClass,
     register_pytest_array_context_factory,
 )
+from boxtree.array_context import PyOpenCLArrayContext as PyOpenCLArrayContextBase
 
 
 if TYPE_CHECKING:

sumpy/distributed.py

@@ -26,7 +26,6 @@
 from typing import TYPE_CHECKING
 
 from boxtree.distributed.calculation import DistributedExpansionWranglerMixin
-
 from pytools import obj_array
 
 from sumpy.fmm import SumpyExpansionWrangler