Merge pull request #19979 from BvB93/memmap

ENH: Add annotations for `np.memmap`

4 files changed, 89 insertions, 12 deletions

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 84a60483d..c78d48cc6 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -196,6 +196,7 @@ from typing import (
     SupportsIndex,
     Final,
     final,
+    ClassVar,
 )
 
 # Ensures that the stubs are picked up
@@ -636,6 +637,17 @@ class _IOProtocol(Protocol):
     def tell(self) -> SupportsIndex: ...
     def seek(self, offset: int, whence: int, /) -> object: ...
 
+# NOTE: `seek`, `write` and `flush` are technically only required
+# for `readwrite`/`write` modes
+class _MemMapIOProtocol(Protocol):
+    def flush(self) -> object: ...
+    def fileno(self) -> SupportsIndex: ...
+    def tell(self) -> int: ...
+    def seek(self, offset: int, whence: int, /) -> object: ...
+    def write(self, s: bytes, /) -> object: ...
+    @property
+    def read(self) -> object: ...
+
 __all__: List[str]
 __path__: List[str]
 __version__: str
@@ -758,18 +770,6 @@ class matrix(ndarray[_ShapeType, _DType_co]):
     def getH(self): ...
     def getI(self): ...
 
-class memmap(ndarray[_ShapeType, _DType_co]):
-    def __new__(
-        subtype,
-        filename: Any,
-        dtype: Any = ...,
-        mode: Any = ...,
-        offset: Any = ...,
-        shape: Any = ...,
-        order: Any = ...,
-    ) -> Any: ...
-    def __getattr__(self, key: str) -> Any: ...
-
 class poly1d:
     def __init__(
         self,
@@ -3828,3 +3828,54 @@ class nditer:
     def shape(self) -> Tuple[int, ...]: ...
     @property
     def value(self) -> Tuple[NDArray[Any], ...]: ...
+
+_MemMapModeKind = L[
+    "readonly", "r",
+    "copyonwrite", "c",
+    "readwrite", "r+",
+    "write", "w+",
+]
+
+class memmap(ndarray[_ShapeType, _DType_co]):
+    __array_priority__: ClassVar[float]
+    filename: str | None
+    offset: int
+    mode: str
+    @overload
+    def __new__(
+        subtype,
+        filename: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _MemMapIOProtocol,
+        dtype: Type[uint8] = ...,
+        mode: _MemMapModeKind = ...,
+        offset: int = ...,
+        shape: None | int | Tuple[int, ...] = ...,
+        order: _OrderKACF = ...,
+    ) -> memmap[Any, dtype[uint8]]: ...
+    @overload
+    def __new__(
+        subtype,
+        filename: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _MemMapIOProtocol,
+        dtype: _DTypeLike[_ScalarType],
+        mode: _MemMapModeKind = ...,
+        offset: int = ...,
+        shape: None | int | Tuple[int, ...] = ...,
+        order: _OrderKACF = ...,
+    ) -> memmap[Any, dtype[_ScalarType]]: ...
+    @overload
+    def __new__(
+        subtype,
+        filename: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _MemMapIOProtocol,
+        dtype: DTypeLike,
+        mode: _MemMapModeKind = ...,
+        offset: int = ...,
+        shape: None | int | Tuple[int, ...] = ...,
+        order: _OrderKACF = ...,
+    ) -> memmap[Any, dtype[Any]]: ...
+    def __array_finalize__(self, obj: memmap[Any, Any]) -> None: ...
+    def __array_wrap__(
+        self,
+        array: memmap[_ShapeType, _DType_co],
+        context: None | Tuple[ufunc, Tuple[Any, ...], int] = ...,
+    ) -> Any: ...
+    def __getitem__(self, index): ...  # TODO
+    def flush(self) -> None: ...
diff --git a/numpy/core/memmap.pyi b/numpy/core/memmap.pyi
new file mode 100644
index 000000000..ba595bf1e
--- /dev/null
+++ b/numpy/core/memmap.pyi
@@ -0,0 +1,5 @@
+from typing import List
+
+from numpy import memmap as memmap
+
+__all__: List[str]
diff --git a/numpy/typing/tests/data/fail/memmap.py b/numpy/typing/tests/data/fail/memmap.py
new file mode 100644
index 000000000..434870b60
--- /dev/null
+++ b/numpy/typing/tests/data/fail/memmap.py
@@ -0,0 +1,5 @@
+import numpy as np
+
+with open("file.txt", "r") as f:
+    np.memmap(f)  # E: No overload variant
+np.memmap("test.txt", shape=[10, 5])  # E: No overload variant
diff --git a/numpy/typing/tests/data/reveal/memmap.py b/numpy/typing/tests/data/reveal/memmap.py
new file mode 100644
index 000000000..c1d8edc67
--- /dev/null
+++ b/numpy/typing/tests/data/reveal/memmap.py
@@ -0,0 +1,16 @@
+import numpy as np
+from typing import Any
+
+memmap_obj: np.memmap[Any, np.dtype[np.str_]]
+
+reveal_type(np.memmap.__array_priority__)  # E: float
+reveal_type(memmap_obj.__array_priority__)  # E: float
+reveal_type(memmap_obj.filename)  # E: Union[builtins.str, None]
+reveal_type(memmap_obj.offset)  # E: int
+reveal_type(memmap_obj.mode)  # E: str
+reveal_type(memmap_obj.flush())  # E: None
+
+reveal_type(np.memmap("file.txt", offset=5))  # E: numpy.memmap[Any, numpy.dtype[{uint8}]]
+reveal_type(np.memmap(b"file.txt", dtype=np.float64, shape=(10, 3)))  # E: numpy.memmap[Any, numpy.dtype[{float64}]]
+with open("file.txt", "rb") as f:
+    reveal_type(np.memmap(f, dtype=float, order="K"))  # E: numpy.memmap[Any, numpy.dtype[Any]]