Merge pull request #142 from jorenham/complete/scipy.fft

complete `scipy.fft`

README.md

@@ -80,7 +80,7 @@ pip install scipy-stubs
 | `cluster`     | ✔️                     | ✔️          | ✔️           | ✔️              | 🌕    |
 | `constants`   | ✔️                     | ✔️          | ✔️           | ✔️              | 🌝    |
 | `datasets`    | ✔️                     | ✔️          | ✔️           | ✔️              | 🌝    |
-| `fft`         | ✔️                     | ✔️          | ✔️           | ✔️              | 🌒    |
+| `fft`         | ✔️                     | ✔️          | ✔️           | ✔️              | 🌔    |
 | `fftpack`     | ✔️                     | ✔️          | ✔️           | ✔️              | 🌔    |
 | `integrate`   | ✔️                     | ✔️          | ✔️           | ✔️              | 🌔    |
 | `interpolate` | ✔️                     | ✔️          | ✔️           | ✔️              | 🌓    |

scipy-stubs/fft/__init__.pyi

@@ -1,46 +1,74 @@
-from ._backend import (
-    register_backend as register_backend,
-    set_backend as set_backend,
-    set_global_backend as set_global_backend,
-    skip_backend as skip_backend,
-)
+from contextlib import _GeneratorContextManager
+
+from ._backend import register_backend, set_backend, set_global_backend, skip_backend
 from ._basic import (
-    fft as fft,
-    fft2 as fft2,
-    fftn as fftn,
-    hfft as hfft,
-    hfft2 as hfft2,
-    hfftn as hfftn,
-    ifft as ifft,
-    ifft2 as ifft2,
-    ifftn as ifftn,
-    ihfft as ihfft,
-    ihfft2 as ihfft2,
-    ihfftn as ihfftn,
-    irfft as irfft,
-    irfft2 as irfft2,
-    irfftn as irfftn,
-    rfft as rfft,
-    rfft2 as rfft2,
-    rfftn as rfftn,
-)
-from ._fftlog import fht as fht, fhtoffset as fhtoffset, ifht as ifht
-from ._helper import (
-    fftfreq as fftfreq,
-    fftshift as fftshift,
-    ifftshift as ifftshift,
-    next_fast_len as next_fast_len,
-    prev_fast_len as prev_fast_len,
-    rfftfreq as rfftfreq,
-)
-from ._pocketfft.helper import get_workers as get_workers, set_workers as set_workers
-from ._realtransforms import (
-    dct as dct,
-    dctn as dctn,
-    dst as dst,
-    dstn as dstn,
-    idct as idct,
-    idctn as idctn,
-    idst as idst,
-    idstn as idstn,
+    fft,
+    fft2,
+    fftn,
+    hfft,
+    hfft2,
+    hfftn,
+    ifft,
+    ifft2,
+    ifftn,
+    ihfft,
+    ihfft2,
+    ihfftn,
+    irfft,
+    irfft2,
+    irfftn,
+    rfft,
+    rfft2,
+    rfftn,
 )
+from ._fftlog import fht, fhtoffset, ifht
+from ._helper import fftfreq, fftshift, ifftshift, next_fast_len, prev_fast_len, rfftfreq
+from ._realtransforms import dct, dctn, dst, dstn, idct, idctn, idst, idstn
+
+__all__ = [
+    "dct",
+    "dctn",
+    "dst",
+    "dstn",
+    "fft",
+    "fft2",
+    "fftfreq",
+    "fftn",
+    "fftshift",
+    "fht",
+    "fhtoffset",
+    "get_workers",
+    "hfft",
+    "hfft2",
+    "hfftn",
+    "idct",
+    "idctn",
+    "idst",
+    "idstn",
+    "ifft",
+    "ifft2",
+    "ifftn",
+    "ifftshift",
+    "ifht",
+    "ihfft",
+    "ihfft2",
+    "ihfftn",
+    "irfft",
+    "irfft2",
+    "irfftn",
+    "next_fast_len",
+    "prev_fast_len",
+    "register_backend",
+    "rfft",
+    "rfft2",
+    "rfftfreq",
+    "rfftn",
+    "set_backend",
+    "set_global_backend",
+    "set_workers",
+    "skip_backend",
+]
+
+# originally defined in `scipy.fft._pocketfft.helper`
+def set_workers(workers: int) -> _GeneratorContextManager[None]: ...
+def get_workers() -> int: ...

scipy-stubs/fft/_backend.pyi

@@ -1,11 +1,23 @@
-from scipy._typing import Untyped
+from collections.abc import Mapping, Sequence
+from typing import Any, ClassVar, Protocol, final, type_check_only
+from typing_extensions import TypeVar
 
-class _ScipyBackend:
-    __ua_domain__: str
+from optype import CanWith
+
+_RT_co = TypeVar("_RT_co", covariant=True, default=Any)
+
+@type_check_only
+class _BaseBackend(Protocol[_RT_co]):
+    __ua_domain__: ClassVar = "numpy.scipy.fft"
     @staticmethod
-    def __ua_function__(method: Untyped, args: Untyped, kwargs: Untyped) -> Untyped: ...
+    def __ua_function__(method: str, args: Sequence[object], kwargs: Mapping[str, object]) -> _RT_co: ...
+
+###
+
+@final
+class _ScipyBackend(_BaseBackend): ...
 
-def set_global_backend(backend: Untyped, coerce: bool = False, only: bool = False, try_last: bool = False) -> None: ...
-def register_backend(backend: Untyped) -> None: ...
-def set_backend(backend: Untyped, coerce: bool = False, only: bool = False) -> Untyped: ...
-def skip_backend(backend: Untyped) -> Untyped: ...
+def set_global_backend(backend: _BaseBackend, coerce: bool = False, only: bool = False, try_last: bool = False) -> None: ...
+def register_backend(backend: _BaseBackend) -> None: ...
+def set_backend(backend: _BaseBackend, coerce: bool = False, only: bool = False) -> CanWith[None]: ...
+def skip_backend(backend: _BaseBackend) -> CanWith[None]: ...