scipy.datasets

README.md

@@ -79,16 +79,13 @@ pip install scipy-stubs
 
 ## Development Progress
 
-According to [basedpyright](https://github.com/DetachHead/basedpyright) (stricter than
-pyright), the "type completeness score" is **42.6%**.
-
 | Package or module                 | Stubs status    |
 |---------------------------------- |---------------- |
 | `scipy.__init__`                  | 3: ready        |
 | `scipy._lib`                      | 2: partial      |
 | `scipy.cluster`                   | 1: skeleton     |
 | `scipy.constants`                 | **4: done**     |
-| `scipy.datasets`                  | 2: partial      |
+| `scipy.datasets`                  | **4: done**     |
 | `scipy.fft`                       | 2: partial      |
 | `scipy.fft._pocketfft`            | 2: partial      |
 | `scipy.fftpack`                   | 2: partial      |

scipy-stubs/datasets/__init__.pyi

@@ -1,3 +1,5 @@
-from ._download_all import download_all as download_all
-from ._fetchers import ascent as ascent, electrocardiogram as electrocardiogram, face as face
-from ._utils import clear_cache as clear_cache
+from ._download_all import download_all
+from ._fetchers import ascent, electrocardiogram, face
+from ._utils import clear_cache
+
+__all__ = ["ascent", "clear_cache", "download_all", "electrocardiogram", "face"]

scipy-stubs/datasets/_download_all.pyi

@@ -1,4 +1,4 @@
-from scipy._typing import Untyped
+from os import PathLike
 
-def download_all(path: Untyped | None = None): ...
-def main(): ...
+def download_all(path: str | PathLike[str] | None = None) -> None: ...
+def main() -> None: ...

scipy-stubs/datasets/_fetchers.pyi

@@ -1,9 +1,16 @@
-from scipy._typing import Untyped
-from ._registry import registry as registry, registry_urls as registry_urls
+from typing import Final, Literal, TypeAlias, overload
+from typing_extensions import LiteralString
 
-data_fetcher: Untyped
+import numpy as np
 
-def fetch_data(dataset_name, data_fetcher=...) -> Untyped: ...
-def ascent() -> Untyped: ...
-def electrocardiogram() -> Untyped: ...
-def face(gray: bool = False) -> Untyped: ...
+# TODO: stub `pooch` (this should be a `pooch.code.Pooch`)
+_DataFetcher: TypeAlias = object
+data_fetcher: Final[_DataFetcher]
+
+def fetch_data(dataset_name: LiteralString, data_fetcher: _DataFetcher = ...) -> LiteralString: ...
+def ascent() -> np.ndarray[tuple[Literal[512], Literal[512]], np.dtype[np.uint8]]: ...
+def electrocardiogram() -> np.ndarray[tuple[Literal[108_000]], np.dtype[np.float64]]: ...
+@overload
+def face(gray: Literal[False] = False) -> np.ndarray[tuple[Literal[768], Literal[1_024], Literal[3]], np.dtype[np.uint8]]: ...
+@overload
+def face(gray: Literal[True]) -> np.ndarray[tuple[Literal[768], Literal[1_024]], np.dtype[np.uint8]]: ...

scipy-stubs/datasets/_registry.pyi

@@ -1,5 +1,6 @@
-from scipy._typing import Untyped
+from typing import Final
+from typing_extensions import LiteralString
 
-registry: Untyped
-registry_urls: Untyped
-method_files_map: Untyped
+registry: Final[dict[LiteralString, LiteralString]]
+registry_urls: Final[dict[LiteralString, LiteralString]]
+method_files_map: Final[dict[LiteralString, list[LiteralString]]]

scipy-stubs/datasets/_utils.pyi

@@ -1,4 +1,12 @@
-from scipy._typing import Untyped
-from ._registry import method_files_map as method_files_map
+from collections.abc import Callable
+from typing import TypeAlias
+from typing_extensions import TypeVar
 
-def clear_cache(datasets: Untyped | None = None): ...
+import numpy as np
+
+_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...], default=tuple[int] | tuple[int, int] | tuple[int, int, int])
+_DT = TypeVar("_DT", bound=np.dtype[np.generic], default=np.dtype[np.float64] | np.dtype[np.uint8])
+
+_AnyDataset: TypeAlias = Callable[[], np.ndarray[_ShapeT, _DT]]
+
+def clear_cache(datasets: list[_AnyDataset] | tuple[_AnyDataset, ...] | None = None) -> None: ...