From 22ae671fd13e6e5091458509f57e8a0996c532e3 Mon Sep 17 00:00:00 2001 From: jorenham Date: Wed, 12 Feb 2025 03:19:30 +0100 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix=20`memmap`=20in=20`numpy.=5F?= =?UTF-8?q?core.memmap`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/__init__.pyi | 60 ++---------------------- src/numpy-stubs/_core/memmap.pyi | 78 +++++++++++++++++++++++++++++++- test/static/accept/memmap.pyi | 10 ++-- 3 files changed, 87 insertions(+), 61 deletions(-) diff --git a/src/numpy-stubs/__init__.pyi b/src/numpy-stubs/__init__.pyi index 1b5ec755..5a00b538 100644 --- a/src/numpy-stubs/__init__.pyi +++ b/src/numpy-stubs/__init__.pyi @@ -95,6 +95,7 @@ from numpy._core.fromnumeric import ( ) from numpy._core.function_base import geomspace, linspace, logspace from numpy._core.getlimits import finfo, iinfo +from numpy._core.memmap import memmap from numpy._core.multiarray import ( arange, array, @@ -500,7 +501,10 @@ __all__ = [ # noqa: RUF022 # _core.__all__ "abs", "acos", "acosh", "asin", "asinh", "atan", "atanh", "atan2", "bitwise_invert", "bitwise_left_shift", "bitwise_right_shift", "concat", "pow", "permute_dims", - "memmap", "sctypeDict", "record", "recarray", + "sctypeDict", "record", "recarray", + + # _core.memmap.__all__ + "memmap", # _core.numeric.__all__ "newaxis", "ndarray", "flatiter", "nditer", "nested_iters", "ufunc", "arange", @@ -906,7 +910,6 @@ _FutureScalar: TypeAlias = L["bytes", "str", "object"] _ByteOrderChar: TypeAlias = L["<", ">", "=", "|"] _CastingKind: TypeAlias = L["no", "equiv", "safe", "same_kind", "unsafe"] _ModeKind: TypeAlias = L["raise", "wrap", "clip"] -_ModeKindMM: TypeAlias = L["readonly", "r", "copyonwrite", "c", "readwrite", "r+", "write", "w+"] _PartitionKind: TypeAlias = L["introselect"] _SortKind: TypeAlias = L[ "Q", "quick", "quicksort", @@ -956,9 +959,6 @@ class _SupportsFileMethods(SupportsFlush, Protocol): def tell(self) -> SupportsIndex: ... def seek(self, offset: int, whence: int, /) -> object: ... -@type_check_only -class _SupportsFileMethodsRW(SupportsWrite[bytes], _SupportsFileMethods, Protocol): ... - @type_check_only class _SupportsItem(Protocol[_T_co]): def item(self, /) -> _T_co: ... @@ -5240,56 +5240,6 @@ class ndindex: def __iter__(self) -> Self: ... def __next__(self) -> _Shape: ... -class memmap(ndarray[_ShapeT_co, _DType_co]): - __array_priority__: ClassVar[float] # pyright: ignore[reportIncompatibleMethodOverride] - - filename: str | None - offset: int - mode: str - - @overload - def __new__( - cls, - filename: StrOrBytesPath | _SupportsFileMethodsRW, - dtype: type[uint8] = ..., - mode: _ModeKindMM = ..., - offset: int = ..., - shape: int | tuple[int, ...] | None = ..., - order: _OrderKACF = ..., - ) -> memmap[Any, dtype[uint8]]: ... - @overload - def __new__( - cls, - filename: StrOrBytesPath | _SupportsFileMethodsRW, - dtype: _DTypeLike[_SCT], - mode: _ModeKindMM = ..., - offset: int = ..., - shape: int | tuple[int, ...] | None = ..., - order: _OrderKACF = ..., - ) -> memmap[Any, dtype[_SCT]]: ... - @overload - def __new__( - cls, - filename: StrOrBytesPath | _SupportsFileMethodsRW, - dtype: DTypeLike, - mode: _ModeKindMM = ..., - offset: int = ..., - shape: int | tuple[int, ...] | None = ..., - order: _OrderKACF = ..., - ) -> memmap[Any, dtype[Any]]: ... - - # - def __array_finalize__(self, obj: object) -> None: ... - def __array_wrap__( - self, - array: memmap[_ShapeT_co, _DType_co], - context: tuple[ufunc, tuple[Any, ...], int] | None = ..., - return_scalar: builtins.bool = ..., - ) -> Any: ... - - # - def flush(self) -> None: ... - class vectorize: __doc__: str | None pyfunc: Callable[..., Any] diff --git a/src/numpy-stubs/_core/memmap.pyi b/src/numpy-stubs/_core/memmap.pyi index 45638a28..b1a0ad91 100644 --- a/src/numpy-stubs/_core/memmap.pyi +++ b/src/numpy-stubs/_core/memmap.pyi @@ -1,3 +1,79 @@ -from numpy import memmap # noqa: ICN003 +from _typeshed import Incomplete, StrOrBytesPath, SupportsWrite +from typing import Any, ClassVar, Final, Generic, Literal as L, Protocol, SupportsIndex, TypeAlias, overload, type_check_only +from typing_extensions import Self, TypeVar + +import numpy as np +import numpy.typing as npt +from numpy import _OrderKACF # noqa: ICN003 +from numpy._typing import _DTypeLike, _ShapeLike __all__ = ["memmap"] + +### + +_ShapeT_co = TypeVar("_ShapeT_co", bound=tuple[int, ...], default=Any, covariant=True) +_DType_co = TypeVar("_DType_co", bound=np.dtype[Any], default=np.dtype[Any], covariant=True) +_ScalarT = TypeVar("_ScalarT", bound=np.generic) + +_Mode: TypeAlias = L["r", "c", "r+", "w+"] +_ToMode: TypeAlias = L["readonly", "r", "copyonwrite", "c", "readwrite", "r+", "write", "w+"] +_ToFileName: TypeAlias = StrOrBytesPath | _SupportsFileMethodsRW + +@type_check_only +class _SupportsFileMethodsRW(SupportsWrite[bytes], Protocol): + def fileno(self, /) -> SupportsIndex: ... + def tell(self, /) -> SupportsIndex: ... + def seek(self, offset: int, whence: int, /) -> object: ... + +### + +class memmap(np.ndarray[_ShapeT_co, _DType_co], Generic[_ShapeT_co, _DType_co]): + __array_priority__: ClassVar[float] = -100.0 # pyright: ignore[reportIncompatibleMethodOverride] + + filename: Final[str | None] + offset: Final[int] + mode: Final[_Mode] + + @overload + def __new__( + cls, + filename: _ToFileName, + dtype: type[np.uint8] = ..., + mode: _ToMode = "r+", + offset: int = 0, + shape: int | tuple[int, ...] | None = None, + order: _OrderKACF = "C", + ) -> memmap[Incomplete, np.dtype[np.uint8]]: ... + @overload + def __new__( + cls, + filename: _ToFileName, + dtype: _DType_co, + mode: _ToMode = "r+", + offset: int = 0, + shape: _ShapeLike | None = None, + order: _OrderKACF = "C", + ) -> memmap[Incomplete, _DType_co]: ... + @overload + def __new__( + cls, + filename: _ToFileName, + dtype: _DTypeLike[_ScalarT], + mode: _ToMode = "r+", + offset: int = 0, + shape: _ShapeLike | None = None, + order: _OrderKACF = "C", + ) -> memmap[Incomplete, np.dtype[_ScalarT]]: ... + @overload + def __new__( + cls, + filename: _ToFileName, + dtype: npt.DTypeLike, + mode: _ToMode = "r+", + offset: int = 0, + shape: int | tuple[int, ...] | None = None, + order: _OrderKACF = "C", + ) -> Self: ... + + # + def flush(self) -> None: ... diff --git a/test/static/accept/memmap.pyi b/test/static/accept/memmap.pyi index fd7e9e4b..b1e55206 100644 --- a/test/static/accept/memmap.pyi +++ b/test/static/accept/memmap.pyi @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, Literal from typing_extensions import assert_type import numpy as np @@ -9,12 +9,12 @@ assert_type(np.memmap.__array_priority__, float) assert_type(memmap_obj.__array_priority__, float) assert_type(memmap_obj.filename, str | None) assert_type(memmap_obj.offset, int) -assert_type(memmap_obj.mode, str) +assert_type(memmap_obj.mode, Literal["r", "c", "r+", "w+"]) assert_type(memmap_obj.flush(), None) assert_type(np.memmap("file.txt", offset=5), np.memmap[Any, np.dtype[np.uint8]]) assert_type(np.memmap(b"file.txt", dtype=np.float64, shape=(10, 3)), np.memmap[Any, np.dtype[np.float64]]) with open("file.txt", "rb") as f: - assert_type(np.memmap(f, dtype=float, order="K"), np.memmap[Any, np.dtype[Any]]) - -assert_type(memmap_obj.__array_finalize__(object()), None) + assert_type(np.memmap(f, dtype=float, order="K"), np.memmap) + assert_type(np.memmap(f, dtype=np.float16), np.memmap[Any, np.dtype[np.float16]]) + assert_type(np.memmap(f, dtype=np.dtypes.StringDType()), np.memmap[Any, np.dtypes.StringDType])