From f56271393956a99e7d9673b8628458fc6ad8753c Mon Sep 17 00:00:00 2001 From: Sebastian Rittau Date: Sat, 29 May 2021 18:01:43 +0200 Subject: [PATCH] Use protocols instead of IO for bz2 (#5499) --- stdlib/bz2.pyi | 94 ++++++++++++++++++++++++++++++++++++++-------- stdlib/tarfile.pyi | 30 ++++++++++++++- 2 files changed, 107 insertions(+), 17 deletions(-) diff --git a/stdlib/bz2.pyi b/stdlib/bz2.pyi index 390024a04184..acbbbbbf507a 100644 --- a/stdlib/bz2.pyi +++ b/stdlib/bz2.pyi @@ -1,22 +1,35 @@ -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 = ..., @@ -24,8 +37,8 @@ def open( ) -> BZ2File: ... @overload def open( - filename: StrOrBytesPath, - mode: _OpenTextMode, + filename: _ReadableFileobj, + mode: _ReadTextMode, compresslevel: int = ..., encoding: Optional[str] = ..., errors: Optional[str] = ..., @@ -33,21 +46,72 @@ def open( ) -> 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: ... diff --git a/stdlib/tarfile.pyi b/stdlib/tarfile.pyi index 16f759e986e6..958ba19da2ea 100644 --- a/stdlib/tarfile.pyi +++ b/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] = ...,