Make BufferedWriter generic over a protocol

python · srittau · May 11, 2025 · May 11, 2025 · May 11, 2025 · May 11, 2025
commit 99fb43ee09cf3e4174a9eb3acf125e954e8b3421
diff --git a/stdlib/@tests/test_cases/check_io.py b/stdlib/@tests/test_cases/check_io.py
@@ -1,9 +1,10 @@
-from _io import BufferedReader
+from _io import BufferedReader, BufferedWriter
 from gzip import GzipFile
 from io import FileIO, RawIOBase, TextIOWrapper
 from typing_extensions import assert_type
 
 BufferedReader(RawIOBase())
+BufferedWriter(RawIOBase())
 
 assert_type(TextIOWrapper(FileIO("")).buffer, FileIO)
 assert_type(TextIOWrapper(FileIO(13)).detach(), FileIO)

diff --git a/stdlib/_io.pyi b/stdlib/_io.pyi
@@ -88,6 +88,7 @@ class BytesIO(BufferedIOBase, _BufferedIOBase, BinaryIO):  # type: ignore[misc]
     def readlines(self, size: int | None = None, /) -> list[bytes]: ...
     def seek(self, pos: int, whence: int = 0, /) -> int: ...
 
+@type_check_only
 class _BufferedReaderStream(Protocol):
     def read(self, n: int = ..., /) -> bytes: ...
     # Optional: def readall(self) -> bytes: ...
@@ -122,9 +123,36 @@ class BufferedReader(BufferedIOBase, _BufferedIOBase, BinaryIO, Generic[_Buffere
     def seek(self, target: int, whence: int = 0, /) -> int: ...
     def truncate(self, pos: int | None = None, /) -> int: ...
 
-class BufferedWriter(BufferedIOBase, _BufferedIOBase, BinaryIO):  # type: ignore[misc]  # incompatible definitions of writelines in the base classes
-    raw: RawIOBase
-    def __init__(self, raw: RawIOBase, buffer_size: int = 8192) -> None: ...
+@type_check_only
+class _BufferedWriterStream(Protocol):
+    def write(self, b: WriteableBuffer, /) -> int: ...
+    def seek(self, pos: int, whence: int, /) -> int: ...
+    def tell(self) -> int: ...
+    def truncate(self, size: int, /) -> int: ...
+    def flush(self) -> object: ...
+    def close(self) -> object: ...
+    @property
+    def closed(self) -> bool: ...
+    def writable(self) -> bool: ...
+    def seekable(self) -> bool: ...
+
+    # The following methods just pass through to the underlying stream. Since
+    # not all streams support them, they are marked as optional here, and will
+    # raise an AttributeError if called on a stream that does not support them.
+
+    # @property
+    # def name(self) -> Any: ...  # Type is inconsistent between the various I/O types.
+    # @property
+    # def mode(self) -> str: ...
+    # def fileno(self) -> int: ...
+    # def isatty(self) -> bool: ...
+
+_BufferedWriterStreamT = TypeVar("_BufferedWriterStreamT", bound=_BufferedWriterStream, default=_BufferedWriterStream)
+
+
+class BufferedWriter(BufferedIOBase, _BufferedIOBase, BinaryIO, Generic[_BufferedWriterStreamT]):  # type: ignore[misc]  # incompatible definitions of writelines in the base classes
+    raw: _BufferedWriterStreamT
+    def __init__(self, raw: _BufferedWriterStreamT, buffer_size: int = 8192) -> None: ...
     def write(self, buffer: ReadableBuffer, /) -> int: ...
     def seek(self, target: int, whence: int = 0, /) -> int: ...
     def truncate(self, pos: int | None = None, /) -> int: ...
@@ -138,8 +166,8 @@ class BufferedRandom(BufferedIOBase, _BufferedIOBase, BinaryIO):  # type: ignore
     def peek(self, size: int = 0, /) -> bytes: ...
     def truncate(self, pos: int | None = None, /) -> int: ...
 
-class BufferedRWPair(BufferedIOBase, _BufferedIOBase, Generic[_BufferedReaderStreamT]):
-    def __init__(self, reader: _BufferedReaderStreamT, writer: RawIOBase, buffer_size: int = 8192, /) -> None: ...
+class BufferedRWPair(BufferedIOBase, _BufferedIOBase, Generic[_BufferedReaderStreamT, _BufferedWriterStreamT]):
+    def __init__(self, reader: _BufferedReaderStreamT, writer: _BufferedWriterStreamT, buffer_size: int = 8192, /) -> None: ...
     def peek(self, size: int = 0, /) -> bytes: ...
 
 class _TextIOBase(_IOBase):