ENH: Use custom file-like protocols instead of `typing.IO`

1 files changed, 32 insertions, 17 deletions

diff --git a/numpy/lib/npyio.pyi b/numpy/lib/npyio.pyi
index 1fa689bbe..edf3daf07 100644
--- a/numpy/lib/npyio.pyi
+++ b/numpy/lib/npyio.pyi
@@ -18,6 +18,7 @@ from typing import (
     Callable,
     Pattern,
     Protocol,
+    Iterable,
 )
 
 from numpy import (
@@ -42,6 +43,8 @@ _T = TypeVar("_T")
 _T_contra = TypeVar("_T_contra", contravariant=True)
 _T_co = TypeVar("_T_co", covariant=True)
 _SCT = TypeVar("_SCT", bound=generic)
+_CharType_co = TypeVar("_CharType_co", str, bytes, covariant=True)
+_CharType_contra = TypeVar("_CharType_contra", str, bytes, contravariant=True)
 
 _DTypeLike = Union[
     Type[_SCT],
@@ -52,6 +55,16 @@ _DTypeLike = Union[
 class _SupportsGetItem(Protocol[_T_contra, _T_co]):
     def __getitem__(self, key: _T_contra) -> _T_co: ...
 
+class _SupportsRead(Protocol[_CharType_co]):
+    def read(self) -> _CharType_co: ...
+
+class _SupportsReadSeek(Protocol[_CharType_co]):
+    def read(self, n: int, /) -> _CharType_co: ...
+    def seek(self, offset: int, whence: int, /) -> object: ...
+
+class _SupportsWrite(Protocol[_CharType_contra]):
+    def write(self, s: _CharType_contra, /) -> object: ...
+
 __all__: List[str]
 
 class BagObj(Generic[_T_co]):
@@ -94,7 +107,7 @@ class NpzFile(Mapping[str, NDArray[Any]]):
 # NOTE: Returns a `NpzFile` if file is a zip file;
 # returns an `ndarray`/`memmap` otherwise
 def load(
-    file: str | bytes | os.PathLike[Any] | IO[bytes],
+    file: str | bytes | os.PathLike[Any] | _SupportsReadSeek[bytes],
     mmap_mode: L[None, "r+", "r", "w+", "c"] = ...,
     allow_pickle: bool = ...,
     fix_imports: bool = ...,
@@ -102,27 +115,29 @@ def load(
 ) -> Any: ...
 
 def save(
-    file: str | os.PathLike[str] | IO[bytes],
+    file: str | os.PathLike[str] | _SupportsWrite[bytes],
     arr: ArrayLike,
     allow_pickle: bool = ...,
     fix_imports: bool = ...,
 ) -> None: ...
 
 def savez(
-    file: str | os.PathLike[str] | IO[bytes],
+    file: str | os.PathLike[str] | _SupportsWrite[bytes],
     *args: ArrayLike,
     **kwds: ArrayLike,
 ) -> None: ...
 
 def savez_compressed(
-    file: str | os.PathLike[str] | IO[bytes],
+    file: str | os.PathLike[str] | _SupportsWrite[bytes],
     *args: ArrayLike,
     **kwds: ArrayLike,
 ) -> None: ...
 
+# File-like objects only have to implement `__iter__` and,
+# optionally, `encoding`
 @overload
 def loadtxt(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
     dtype: None = ...,
     comments: str | Sequence[str] = ...,
     delimiter: None | str = ...,
@@ -138,7 +153,7 @@ def loadtxt(
 ) -> NDArray[float64]: ...
 @overload
 def loadtxt(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
     dtype: _DTypeLike[_SCT],
     comments: str | Sequence[str] = ...,
     delimiter: None | str = ...,
@@ -154,7 +169,7 @@ def loadtxt(
 ) -> NDArray[_SCT]: ...
 @overload
 def loadtxt(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
     dtype: DTypeLike,
     comments: str | Sequence[str] = ...,
     delimiter: None | str = ...,
@@ -170,7 +185,7 @@ def loadtxt(
 ) -> NDArray[Any]: ...
 
 def savetxt(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | _SupportsWrite[str] | _SupportsWrite[bytes],
     X: ArrayLike,
     fmt: str | Sequence[str] = ...,
     delimiter: str = ...,
@@ -183,14 +198,14 @@ def savetxt(
 
 @overload
 def fromregex(
-    file: str | os.PathLike[str] | IO[Any],
+    file: str | os.PathLike[str] | _SupportsRead[str] | _SupportsRead[bytes],
     regexp: str | bytes | Pattern[Any],
     dtype: _DTypeLike[_SCT],
     encoding: None | str = ...
 ) -> NDArray[_SCT]: ...
 @overload
 def fromregex(
-    file: str | os.PathLike[str] | IO[Any],
+    file: str | os.PathLike[str] | _SupportsRead[str] | _SupportsRead[bytes],
     regexp: str | bytes | Pattern[Any],
     dtype: DTypeLike,
     encoding: None | str = ...
@@ -199,21 +214,21 @@ def fromregex(
 # TODO: Sort out arguments
 @overload
 def genfromtxt(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
     dtype: None = ...,
     *args: Any,
     **kwargs: Any,
 ) -> NDArray[float64]: ...
 @overload
 def genfromtxt(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
     dtype: _DTypeLike[_SCT],
     *args: Any,
     **kwargs: Any,
 ) -> NDArray[_SCT]: ...
 @overload
 def genfromtxt(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
     dtype: DTypeLike,
     *args: Any,
     **kwargs: Any,
@@ -221,14 +236,14 @@ def genfromtxt(
 
 @overload
 def recfromtxt(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
     *,
     usemask: L[False] = ...,
     **kwargs: Any,
 ) -> recarray[Any, dtype[void]]: ...
 @overload
 def recfromtxt(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
     *,
     usemask: L[True],
     **kwargs: Any,
@@ -236,14 +251,14 @@ def recfromtxt(
 
 @overload
 def recfromcsv(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
     *,
     usemask: L[False] = ...,
     **kwargs: Any,
 ) -> recarray[Any, dtype[void]]: ...
 @overload
 def recfromcsv(
-    fname: str | os.PathLike[str] | IO[Any],
+    fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
     *,
     usemask: L[True],
     **kwargs: Any,