From 9ab30c98df6ef01daaa223b5d1aff6d239fe9fb3 Mon Sep 17 00:00:00 2001 From: jorenham Date: Thu, 6 Feb 2025 17:56:25 +0100 Subject: [PATCH] =?UTF-8?q?=F0=9F=8F=B7=EF=B8=8F=20stub=20`=5Fdatasource`?= =?UTF-8?q?=20and=20fix=20`=5Fnpyio=5Fimpl`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/lib/_datasource.pyi | 30 ++++ src/numpy-stubs/lib/_npyio_impl.pyi | 247 ++++++++++++++-------------- src/numpy-stubs/lib/npyio.pyi | 1 + test/static/accept/npyio.pyi | 4 +- 4 files changed, 155 insertions(+), 127 deletions(-) create mode 100644 src/numpy-stubs/lib/_datasource.pyi diff --git a/src/numpy-stubs/lib/_datasource.pyi b/src/numpy-stubs/lib/_datasource.pyi new file mode 100644 index 0000000..ad52b7f --- /dev/null +++ b/src/numpy-stubs/lib/_datasource.pyi @@ -0,0 +1,30 @@ +from _typeshed import OpenBinaryMode, OpenTextMode +from pathlib import Path +from typing import IO, Any, TypeAlias + +_Mode: TypeAlias = OpenBinaryMode | OpenTextMode + +### + +# exported in numpy.lib.nppyio +class DataSource: + def __init__(self, /, destpath: Path | str | None = ...) -> None: ... + def __del__(self, /) -> None: ... + def abspath(self, /, path: str) -> str: ... + def exists(self, /, path: str) -> bool: ... + + # Whether the file-object is opened in string or bytes mode (by default) + # depends on the file-extension of `path` + def open(self, /, path: str, mode: _Mode = "r", encoding: str | None = None, newline: str | None = None) -> IO[Any]: ... + +class Repository(DataSource): + def __init__(self, /, baseurl: str, destpath: str | None = ...) -> None: ... + def listdir(self, /) -> list[str]: ... + +def open( + path: str, + mode: _Mode = "r", + destpath: str | None = ..., + encoding: str | None = None, + newline: str | None = None, +) -> IO[Any]: ... diff --git a/src/numpy-stubs/lib/_npyio_impl.pyi b/src/numpy-stubs/lib/_npyio_impl.pyi index 6cbe867..e5d823f 100644 --- a/src/numpy-stubs/lib/_npyio_impl.pyi +++ b/src/numpy-stubs/lib/_npyio_impl.pyi @@ -3,14 +3,16 @@ import zipfile from _typeshed import StrOrBytesPath, StrPath, SupportsKeysAndGetItem, SupportsRead, SupportsWrite from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence from re import Pattern -from typing import IO, Any, Generic, Literal as L, Protocol, overload, type_check_only -from typing_extensions import Self, TypeVar, deprecated +from typing import IO, Any, ClassVar, Generic, Literal as L, Protocol, TypeAlias, overload, type_check_only +from typing_extensions import Self, TypeVar, deprecated, override import numpy as np from numpy._core.multiarray import packbits, unpackbits from numpy._typing import ArrayLike, DTypeLike, NDArray, _DTypeLike, _SupportsArrayFunc from numpy.ma.mrecords import MaskedRecords +from ._datasource import DataSource as DataSource + __all__ = [ "fromregex", "genfromtxt", @@ -24,164 +26,175 @@ __all__ = [ "unpackbits", ] -_T = TypeVar("_T") _T_co = TypeVar("_T_co", covariant=True) _SCT = TypeVar("_SCT", bound=np.generic) +_SCT_co = TypeVar("_SCT_co", bound=np.generic, default=Any, covariant=True) + +_FName: TypeAlias = StrPath | Iterable[str] | Iterable[bytes] +_FNameRead: TypeAlias = StrPath | SupportsRead[str] | SupportsRead[bytes] +_FNameWriteBytes: TypeAlias = StrPath | SupportsWrite[bytes] +_FNameWrite: TypeAlias = _FNameWriteBytes | SupportsWrite[bytes] @type_check_only class _SupportsReadSeek(SupportsRead[_T_co], Protocol[_T_co]): def seek(self, offset: int, whence: int, /) -> object: ... +### + +# undocumented class BagObj(Generic[_T_co]): - def __init__(self, obj: SupportsKeysAndGetItem[str, _T_co]) -> None: ... - def __getattribute__(self, key: str) -> _T_co: ... + def __init__(self, /, obj: SupportsKeysAndGetItem[str, _T_co]) -> None: ... + def __getattribute__(self, key: str, /) -> _T_co: ... def __dir__(self) -> list[str]: ... -class NpzFile(Mapping[str, NDArray[Any]]): - zip: zipfile.ZipFile +# exported in numpy.lib.nppyio +class NpzFile(Mapping[str, NDArray[_SCT_co]]): + _MAX_REPR_ARRAY_COUNT: ClassVar[int] = 5 + + zip: zipfile.ZipFile | None fid: IO[str] | None files: list[str] allow_pickle: bool pickle_kwargs: Mapping[str, Any] | None - _MAX_REPR_ARRAY_COUNT: int - # Represent `f` as a mutable property so we can access the type of `self` - @property - def f(self: _T) -> BagObj[_T]: ... - @f.setter - def f(self: _T, value: BagObj[_T], /) -> None: ... + f: BagObj[NpzFile[_SCT_co]] | None + + # def __init__( self, - fid: IO[str], - own_fid: bool = ..., - allow_pickle: bool = ..., - pickle_kwargs: Mapping[str, object] | None = ..., + /, + fid: IO[Any], + own_fid: bool = False, + allow_pickle: bool = False, + pickle_kwargs: Mapping[str, object] | None = None, + *, + max_header_size: int = 10_000, ) -> None: ... + def __del__(self) -> None: ... def __enter__(self) -> Self: ... def __exit__(self, cls: type[BaseException] | None, e: BaseException | None, tb: types.TracebackType | None, /) -> None: ... - def close(self) -> None: ... - def __del__(self) -> None: ... - def __iter__(self) -> Iterator[str]: ... + @override def __len__(self) -> int: ... - def __getitem__(self, key: str, /) -> NDArray[Any]: ... - def __contains__(self, key: str, /) -> bool: ... - -class DataSource: - def __init__(self, destpath: StrPath | None = ...) -> None: ... - def __del__(self) -> None: ... - def abspath(self, path: str) -> str: ... - def exists(self, path: str) -> bool: ... - - # Whether the file-object is opened in string or bytes mode (by default) - # depends on the file-extension of `path` - def open(self, path: str, mode: str = ..., encoding: str | None = ..., newline: str | None = ...) -> IO[Any]: ... + @override + def __iter__(self) -> Iterator[str]: ... + @override + def __getitem__(self, key: str, /) -> NDArray[_SCT_co]: ... + def close(self) -> None: ... -# NOTE: Returns a `NpzFile` if file is a zip file; -# returns an `ndarray`/`memmap` otherwise +# NOTE: Returns `NpzFile` if file is a zip file; returns `ndarray`/`memmap` otherwise def load( file: StrOrBytesPath | _SupportsReadSeek[bytes], - mmap_mode: L["r+", "r", "w+", "c"] | None = ..., - allow_pickle: bool = ..., - fix_imports: bool = ..., - encoding: L["ASCII", "latin1", "bytes"] = ..., + mmap_mode: L["r+", "r", "w+", "c"] | None = None, + allow_pickle: bool = False, + fix_imports: bool = True, + encoding: L["ASCII", "latin1", "bytes"] = "ASCII", + *, + max_header_size: int = 10_000, ) -> Any: ... + +# @overload -def save(file: StrPath | SupportsWrite[bytes], arr: ArrayLike, allow_pickle: bool = ...) -> None: ... +def save(file: _FNameWriteBytes, arr: ArrayLike, allow_pickle: bool = True) -> None: ... @overload @deprecated("The 'fix_imports' flag is deprecated in NumPy 2.1.") -def save(file: StrPath | SupportsWrite[bytes], arr: ArrayLike, allow_pickle: bool = ..., *, fix_imports: bool) -> None: ... +def save(file: _FNameWriteBytes, arr: ArrayLike, allow_pickle: bool, fix_imports: bool) -> None: ... @overload @deprecated("The 'fix_imports' flag is deprecated in NumPy 2.1.") -def save(file: StrPath | SupportsWrite[bytes], arr: ArrayLike, allow_pickle: bool, fix_imports: bool) -> None: ... -def savez(file: StrPath | SupportsWrite[bytes], *args: ArrayLike, allow_pickle: bool = ..., **kwds: ArrayLike) -> None: ... -def savez_compressed( - file: StrPath | SupportsWrite[bytes], - *args: ArrayLike, - allow_pickle: bool = ..., - **kwds: ArrayLike, -) -> None: ... +def save(file: _FNameWriteBytes, arr: ArrayLike, allow_pickle: bool = True, *, fix_imports: bool) -> None: ... + +# +def savez(file: _FNameWriteBytes, *args: ArrayLike, allow_pickle: bool = True, **kwds: ArrayLike) -> None: ... + +# +def savez_compressed(file: _FNameWriteBytes, *args: ArrayLike, allow_pickle: bool = True, **kwds: ArrayLike) -> None: ... # File-like objects only have to implement `__iter__` and, # optionally, `encoding` @overload def loadtxt( - fname: StrPath | Iterable[str] | Iterable[bytes], - dtype: None = ..., - comments: str | Sequence[str] | None = ..., - delimiter: str | None = ..., - converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = ..., - skiprows: int = ..., - usecols: int | Sequence[int] | None = ..., - unpack: bool = ..., - ndmin: L[0, 1, 2] = ..., - encoding: str | None = ..., - max_rows: int | None = ..., + fname: _FName, + dtype: None = None, + comments: str | Sequence[str] | None = "#", + delimiter: str | None = None, + converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = None, + skiprows: int = 0, + usecols: int | Sequence[int] | None = None, + unpack: bool = False, + ndmin: L[0, 1, 2] = 0, + encoding: str | None = None, + max_rows: int | None = None, *, - quotechar: str | None = ..., - like: _SupportsArrayFunc | None = ..., + quotechar: str | None = None, + like: _SupportsArrayFunc | None = None, ) -> NDArray[np.float64]: ... @overload def loadtxt( - fname: StrPath | Iterable[str] | Iterable[bytes], + fname: _FName, dtype: _DTypeLike[_SCT], - comments: str | Sequence[str] | None = ..., - delimiter: str | None = ..., - converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = ..., - skiprows: int = ..., - usecols: int | Sequence[int] | None = ..., - unpack: bool = ..., - ndmin: L[0, 1, 2] = ..., - encoding: str | None = ..., - max_rows: int | None = ..., + comments: str | Sequence[str] | None = "#", + delimiter: str | None = None, + converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = None, + skiprows: int = 0, + usecols: int | Sequence[int] | None = None, + unpack: bool = False, + ndmin: L[0, 1, 2] = 0, + encoding: str | None = None, + max_rows: int | None = None, *, - quotechar: str | None = ..., - like: _SupportsArrayFunc | None = ..., + quotechar: str | None = None, + like: _SupportsArrayFunc | None = None, ) -> NDArray[_SCT]: ... @overload def loadtxt( - fname: StrPath | Iterable[str] | Iterable[bytes], + fname: _FName, dtype: DTypeLike, - comments: str | Sequence[str] | None = ..., - delimiter: str | None = ..., - converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = ..., - skiprows: int = ..., - usecols: int | Sequence[int] | None = ..., - unpack: bool = ..., - ndmin: L[0, 1, 2] = ..., - encoding: str | None = ..., - max_rows: int | None = ..., + comments: str | Sequence[str] | None = "#", + delimiter: str | None = None, + converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = None, + skiprows: int = 0, + usecols: int | Sequence[int] | None = None, + unpack: bool = False, + ndmin: L[0, 1, 2] = 0, + encoding: str | None = None, + max_rows: int | None = None, *, - quotechar: str | None = ..., - like: _SupportsArrayFunc | None = ..., + quotechar: str | None = None, + like: _SupportsArrayFunc | None = None, ) -> NDArray[Any]: ... + +# def savetxt( - fname: StrPath | SupportsWrite[str] | SupportsWrite[bytes], + fname: StrPath | _FNameWrite, X: ArrayLike, - fmt: str | Sequence[str] = ..., - delimiter: str = ..., - newline: str = ..., - header: str = ..., - footer: str = ..., - comments: str = ..., - encoding: str | None = ..., + fmt: str | Sequence[str] = "%.18e", + delimiter: str = " ", + newline: str = "\n", + header: str = "", + footer: str = "", + comments: str = "# ", + encoding: str | None = None, ) -> None: ... + +# @overload def fromregex( - file: StrPath | SupportsRead[str] | SupportsRead[bytes], + file: _FNameRead, regexp: str | bytes | Pattern[Any], dtype: _DTypeLike[_SCT], - encoding: str | None = ..., + encoding: str | None = None, ) -> NDArray[_SCT]: ... @overload def fromregex( - file: StrPath | SupportsRead[str] | SupportsRead[bytes], + file: _FNameRead, regexp: str | bytes | Pattern[Any], dtype: DTypeLike, - encoding: str | None = ..., + encoding: str | None = None, ) -> NDArray[Any]: ... + +# @overload def genfromtxt( - fname: StrPath | Iterable[str] | Iterable[bytes], - dtype: None = ..., + fname: _FName, + dtype: None = None, comments: str = ..., delimiter: str | int | Iterable[int] | None = ..., skip_header: int = ..., @@ -209,7 +222,7 @@ def genfromtxt( ) -> NDArray[Any]: ... @overload def genfromtxt( - fname: StrPath | Iterable[str] | Iterable[bytes], + fname: _FName, dtype: _DTypeLike[_SCT], comments: str = ..., delimiter: str | int | Iterable[int] | None = ..., @@ -238,7 +251,7 @@ def genfromtxt( ) -> NDArray[_SCT]: ... @overload def genfromtxt( - fname: StrPath | Iterable[str] | Iterable[bytes], + fname: _FName, dtype: DTypeLike, comments: str = ..., delimiter: str | int | Iterable[int] | None = ..., @@ -265,31 +278,15 @@ def genfromtxt( ndmin: L[0, 1, 2] = ..., like: _SupportsArrayFunc | None = ..., ) -> NDArray[Any]: ... + +# @overload -def recfromtxt( - fname: StrPath | Iterable[str] | Iterable[bytes], - *, - usemask: L[False] = ..., - **kwargs: object, -) -> np.recarray[Any, np.dtype[np.record]]: ... +def recfromtxt(fname: _FName, *, usemask: L[False] = False, **kwargs: object) -> np.recarray[Any, np.dtype[np.record]]: ... @overload -def recfromtxt( - fname: StrPath | Iterable[str] | Iterable[bytes], - *, - usemask: L[True], - **kwargs: object, -) -> MaskedRecords[Any, np.dtype[np.void]]: ... +def recfromtxt(fname: _FName, *, usemask: L[True], **kwargs: object) -> MaskedRecords[Any, np.dtype[np.void]]: ... + +# @overload -def recfromcsv( - fname: StrPath | Iterable[str] | Iterable[bytes], - *, - usemask: L[False] = ..., - **kwargs: object, -) -> np.recarray[Any, np.dtype[np.record]]: ... +def recfromcsv(fname: _FName, *, usemask: L[False] = False, **kwargs: object) -> np.recarray[Any, np.dtype[np.record]]: ... @overload -def recfromcsv( - fname: StrPath | Iterable[str] | Iterable[bytes], - *, - usemask: L[True], - **kwargs: object, -) -> MaskedRecords[Any, np.dtype[np.void]]: ... +def recfromcsv(fname: _FName, *, usemask: L[True], **kwargs: object) -> MaskedRecords[Any, np.dtype[np.void]]: ... diff --git a/src/numpy-stubs/lib/npyio.pyi b/src/numpy-stubs/lib/npyio.pyi index c3258e8..fd3ae8f 100644 --- a/src/numpy-stubs/lib/npyio.pyi +++ b/src/numpy-stubs/lib/npyio.pyi @@ -1,4 +1,5 @@ from numpy.lib._npyio_impl import ( DataSource as DataSource, NpzFile as NpzFile, + __doc__ as __doc__, ) diff --git a/test/static/accept/npyio.pyi b/test/static/accept/npyio.pyi index ceb3cfe..103e9e2 100644 --- a/test/static/accept/npyio.pyi +++ b/test/static/accept/npyio.pyi @@ -29,12 +29,12 @@ class BytesReader: bytes_writer: BytesWriter bytes_reader: BytesReader -assert_type(npz_file.zip, zipfile.ZipFile) +assert_type(npz_file.zip, zipfile.ZipFile | None) assert_type(npz_file.fid, IO[str] | None) assert_type(npz_file.files, list[str]) assert_type(npz_file.allow_pickle, bool) assert_type(npz_file.pickle_kwargs, Mapping[str, Any] | None) -assert_type(npz_file.f, BagObj[np.lib.npyio.NpzFile]) +assert_type(npz_file.f, BagObj[np.lib.npyio.NpzFile] | None) assert_type(npz_file["test"], npt.NDArray[Any]) assert_type(len(npz_file), int) with npz_file as f: