Add type annotations around sumpy.fmm - #300
Conversation
| if self.previous_flag is not None or self.previous_no_cache_kernels is not None: | ||
| raise RuntimeError("cannot reuse the 'CacheMode' context manager") |
There was a problem hiding this comment.
pyright was complaining about setting attributes outside of __init__, so I reworked this a bit. It now yells if it's nested like with ... as cm: with cm: ...
| import numpy | ||
| from numpy.typing import DTypeLike | ||
| from optype.numpy import Array2D | ||
| from pyvkfft.opencl import VkFFTApp |
There was a problem hiding this comment.
Is it ok to make this a dependency for typing? It's mostly just used in some return types.
There was a problem hiding this comment.
Oh sure! Just make sure it's in whatever uv uses for dev dependencies, so that it's not pulled in at runtime.
|
|
||
| @memoize_method | ||
| def multipole_expansion(self, order: int): | ||
| def multipole_expansion(self, order: int) -> MultipoleExpansionBase: |
There was a problem hiding this comment.
I'm not sure about returning the base classes here. It makes sense, but they're currently not implemented very nicely for something like this (e.g. pyright complains MultipoleExpansionBase is not callable because it has no __call__ stub).
| actx: ArrayContext, | ||
| fft_app: tuple[Any, FFTBackend], | ||
| input_vec: Any, | ||
| fft_app: tuple[lp.TranslationUnit | VkFFTApp, FFTBackend], |
There was a problem hiding this comment.
Not a big fan of this being a tuple (the first and second element could get out of sync?). Any reason to not just do an isinstance?
| def reorder_potentials( | ||
| self, potentials: Array | obj_array.ObjectArray1D[Array] | ||
| ) -> Array | obj_array.ObjectArray1D[Array]: |
There was a problem hiding this comment.
This should also be added to boxtree. That currently just takes Array.
|
|
||
| def finalize_potentials(self, actx: ArrayContext, potentials): | ||
| @override | ||
| def finalize_potentials(self, actx: ArrayContext, potentials: Array) -> Array: |
There was a problem hiding this comment.
This should probably also be a Array | ObjectArray, but needs a little update in boxtree as well.
|
Thanks for sifting through this! |
xref: inducer/boxtree#125