🐛 fix memmap in numpy._core.memmap

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]

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: ...

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])