Use protocols instead of IO for bz2 (#5499)

Author: srittau

stdlib/bz2.pyi

@@ -1,53 +1,117 @@
-import io
+import _compression
 import sys
+from _compression import BaseStream
 from _typeshed import ReadableBuffer, StrOrBytesPath, WriteableBuffer
-from typing import IO, Any, Iterable, List, Optional, TextIO, TypeVar, Union, overload
+from typing import IO, Any, Iterable, List, Optional, Protocol, TextIO, TypeVar, Union, overload
 from typing_extensions import Literal, SupportsIndex
 
-_PathOrFile = Union[StrOrBytesPath, IO[bytes]]
+# The following attributes and methods are optional:
+# def fileno(self) -> int: ...
+# def close(self) -> object: ...
+class _ReadableFileobj(_compression._Reader, Protocol): ...
+
+class _WritableFileobj(Protocol):
+    def write(self, __b: bytes) -> object: ...
+    # The following attributes and methods are optional:
+    # def fileno(self) -> int: ...
+    # def close(self) -> object: ...
+
 _T = TypeVar("_T")
 
 def compress(data: bytes, compresslevel: int = ...) -> bytes: ...
 def decompress(data: bytes) -> bytes: ...
 
-_OpenBinaryMode = Literal["r", "rb", "w", "wb", "x", "xb", "a", "ab"]
-_OpenTextMode = Literal["rt", "wt", "xt", "at"]
+_ReadBinaryMode = Literal["", "r", "rb"]
+_WriteBinaryMode = Literal["w", "wb", "x", "xb", "a", "ab"]
+_ReadTextMode = Literal["rt"]
+_WriteTextMode = Literal["wt", "xt", "at"]
 
 @overload
 def open(
-    filename: _PathOrFile,
-    mode: _OpenBinaryMode = ...,
+    filename: _ReadableFileobj,
+    mode: _ReadBinaryMode = ...,
     compresslevel: int = ...,
     encoding: None = ...,
     errors: None = ...,
     newline: None = ...,
 ) -> BZ2File: ...
 @overload
 def open(
-    filename: StrOrBytesPath,
-    mode: _OpenTextMode,
+    filename: _ReadableFileobj,
+    mode: _ReadTextMode,
     compresslevel: int = ...,
     encoding: Optional[str] = ...,
     errors: Optional[str] = ...,
     newline: Optional[str] = ...,
 ) -> TextIO: ...
 @overload
 def open(
-    filename: _PathOrFile,
-    mode: str,
+    filename: _WritableFileobj,
+    mode: _WriteBinaryMode,
+    compresslevel: int = ...,
+    encoding: None = ...,
+    errors: None = ...,
+    newline: None = ...,
+) -> BZ2File: ...
+@overload
+def open(
+    filename: _WritableFileobj,
+    mode: _WriteTextMode,
     compresslevel: int = ...,
     encoding: Optional[str] = ...,
     errors: Optional[str] = ...,
     newline: Optional[str] = ...,
-) -> Union[BZ2File, TextIO]: ...
+) -> TextIO: ...
+@overload
+def open(
+    filename: StrOrBytesPath,
+    mode: Union[_ReadBinaryMode, _WriteBinaryMode] = ...,
+    compresslevel: int = ...,
+    encoding: None = ...,
+    errors: None = ...,
+    newline: None = ...,
+) -> BZ2File: ...
+@overload
+def open(
+    filename: StrOrBytesPath,
+    mode: Union[_ReadTextMode, _WriteTextMode],
+    compresslevel: int = ...,
+    encoding: Optional[str] = ...,
+    errors: Optional[str] = ...,
+    newline: Optional[str] = ...,
+) -> TextIO: ...
 
-class BZ2File(io.BufferedIOBase, IO[bytes]):
+class BZ2File(BaseStream, IO[bytes]):
     def __enter__(self: _T) -> _T: ...
     if sys.version_info >= (3, 9):
-        def __init__(self, filename: _PathOrFile, mode: str = ..., *, compresslevel: int = ...) -> None: ...
+        @overload
+        def __init__(self, filename: _WritableFileobj, mode: _WriteBinaryMode, *, compresslevel: int = ...) -> None: ...
+        @overload
+        def __init__(self, filename: _ReadableFileobj, mode: _ReadBinaryMode = ..., *, compresslevel: int = ...) -> None: ...
+        @overload
+        def __init__(
+            self, filename: StrOrBytesPath, mode: Union[_ReadBinaryMode, _WriteBinaryMode] = ..., *, compresslevel: int = ...
+        ) -> None: ...
     else:
+        @overload
+        def __init__(
+            self, filename: _WritableFileobj, mode: _WriteBinaryMode, buffering: Optional[Any] = ..., compresslevel: int = ...
+        ) -> None: ...
+        @overload
+        def __init__(
+            self,
+            filename: _ReadableFileobj,
+            mode: _ReadBinaryMode = ...,
+            buffering: Optional[Any] = ...,
+            compresslevel: int = ...,
+        ) -> None: ...
+        @overload
         def __init__(
-            self, filename: _PathOrFile, mode: str = ..., buffering: Optional[Any] = ..., compresslevel: int = ...
+            self,
+            filename: StrOrBytesPath,
+            mode: Union[_ReadBinaryMode, _WriteBinaryMode] = ...,
+            buffering: Optional[Any] = ...,
+            compresslevel: int = ...,
         ) -> None: ...
     def read(self, size: Optional[int] = ...) -> bytes: ...
     def read1(self, size: int = ...) -> bytes: ...

stdlib/tarfile.pyi

@@ -1,3 +1,4 @@
+import bz2
 import io
 import sys
 from _typeshed import StrOrBytesPath, StrPath
@@ -17,6 +18,12 @@ class _Fileobj(Protocol):
     # name: str | bytes
     # mode: Literal["rb", "r+b", "wb", "xb"]
 
+class _Bz2ReadableFileobj(bz2._ReadableFileobj):
+    def close(self) -> object: ...
+
+class _Bz2WritableFileobj(bz2._WritableFileobj):
+    def close(self) -> object: ...
+
 # tar constants
 NUL: bytes
 BLOCKSIZE: int
@@ -191,12 +198,31 @@ class TarFile:
         debug: Optional[int] = ...,
         errorlevel: Optional[int] = ...,
     ) -> TarFile: ...
+    @overload
     @classmethod
     def bz2open(
         cls,
         name: Optional[StrOrBytesPath],
-        mode: Literal["r", "w", "x"] = ...,
-        fileobj: Optional[IO[bytes]] = ...,
+        mode: Literal["w", "x"],
+        fileobj: Optional[_Bz2WritableFileobj] = ...,
+        compresslevel: int = ...,
+        *,
+        format: Optional[int] = ...,
+        tarinfo: Optional[Type[TarInfo]] = ...,
+        dereference: Optional[bool] = ...,
+        ignore_zeros: Optional[bool] = ...,
+        encoding: Optional[str] = ...,
+        pax_headers: Optional[Mapping[str, str]] = ...,
+        debug: Optional[int] = ...,
+        errorlevel: Optional[int] = ...,
+    ) -> TarFile: ...
+    @overload
+    @classmethod
+    def bz2open(
+        cls,
+        name: Optional[StrOrBytesPath],
+        mode: Literal["r"] = ...,
+        fileobj: Optional[_Bz2ReadableFileobj] = ...,
         compresslevel: int = ...,
         *,
         format: Optional[int] = ...,