Accurate return types for ZipFile.open() and zipfile.Path.open()

stdlib/zipfile/__init__.pyi

@@ -1,6 +1,6 @@
 import io
 import sys
-from _typeshed import SizedBuffer, StrOrBytesPath, StrPath
+from _typeshed import ReadableBuffer, SizedBuffer, StrOrBytesPath, StrPath
 from collections.abc import Callable, Iterable, Iterator
 from io import TextIOWrapper
 from os import PathLike
@@ -30,7 +30,6 @@ _DateTuple = tuple[int, int, int, int, int, int]  # noqa: Y026
 _ZipFileMode = Literal["r", "w", "x", "a"]  # noqa: Y026
 
 _ReadWriteMode: TypeAlias = Literal["r", "w"]
-_ReadWriteBinaryMode: TypeAlias = Literal["r", "w", "rb", "wb"]
 
 class BadZipFile(Exception): ...
 
@@ -91,6 +90,12 @@ class ZipExtFile(io.BufferedIOBase):
     def read1(self, n: int | None) -> bytes: ...  # type: ignore[override]
     def seek(self, offset: int, whence: int = 0) -> int: ...
 
+class _ZipWriteFile(io.BufferedIOBase):  # undocumented
+    def __init__(self, zf: ZipFile, zinfo: ZipInfo, zip64: bool) -> None: ...
+    def writable(self) -> Literal[True]: ...
+    def write(self, data: ReadableBuffer) -> int: ...
+    def close(self) -> None: ...
+
 class _Writer(Protocol):
     def write(self, s: str, /) -> object: ...
 
@@ -225,8 +230,17 @@ class ZipFile:
     def getinfo(self, name: str) -> ZipInfo: ...
     def infolist(self) -> list[ZipInfo]: ...
     def namelist(self) -> list[str]: ...
+    @overload
+    def open(
+        self, name: str | ZipInfo, mode: Literal["r"] = "r", pwd: bytes | None = None, *, force_zip64: bool = False
+    ) -> ZipExtFile: ...
+    @overload
+    def open(
+        self, name: str | ZipInfo, mode: Literal["w"] = ..., pwd: bytes | None = None, *, force_zip64: bool = False
+    ) -> _ZipWriteFile: ...
+    @overload
     def open(
-        self, name: str | ZipInfo, mode: _ReadWriteMode = "r", pwd: bytes | None = None, *, force_zip64: bool = False
+        self, name: str | ZipInfo, mode: _ReadWriteMode, pwd: bytes | None = None, *, force_zip64: bool = False
     ) -> IO[bytes]: ...
     def extract(self, member: str | ZipInfo, path: StrPath | None = None, pwd: bytes | None = None) -> str: ...
     def extractall(
@@ -335,11 +349,16 @@ else:
                 pwd: bytes | None = None,
             ) -> TextIOWrapper: ...
             @overload
-            def open(self, mode: Literal["rb", "wb"], *, pwd: bytes | None = None) -> IO[bytes]: ...
+            def open(self, mode: Literal["rb"], *, pwd: bytes | None = None) -> ZipExtFile: ...
+            @overload
+            def open(self, mode: Literal["wb"], *, pwd: bytes | None = None) -> _ZipWriteFile: ...
         else:
-            def open(
-                self, mode: _ReadWriteBinaryMode = "r", pwd: bytes | None = None, *, force_zip64: bool = False
-            ) -> IO[bytes]: ...
+            @overload
+            def open(self, mode: Literal["r"] = "r", pwd: bytes | None = None, *, force_zip64: bool = False) -> ZipExtFile: ...
+            @overload
+            def open(self, mode: Literal["w"] = ..., pwd: bytes | None = None, *, force_zip64: bool = False) -> _ZipWriteFile: ...
+            @overload
+            def open(self, mode: _ReadWriteMode, pwd: bytes | None = None, *, force_zip64: bool = False) -> IO[bytes]: ...
 
         if sys.version_info >= (3, 10):
             def iterdir(self) -> Iterator[Self]: ...

stdlib/zipfile/_path/__init__.pyi

@@ -58,7 +58,9 @@
             pwd: bytes | None = None,
         ) -> TextIOWrapper: ...
         @overload
-        def open(self, mode: Literal["rb", "wb"], *, pwd: bytes | None = None) -> IO[bytes]: ...
+        def open(self, mode: Literal["rb"], *, pwd: bytes | None = None) -> ZipExtFile: ...
+        @overload
+        def open(self, mode: Literal["wb"], *, pwd: bytes | None = None) -> _ZipWriteFile: ...
         def iterdir(self) -> Iterator[Self]: ...
         def is_dir(self) -> bool: ...
         def is_file(self) -> bool: ...