Turn TextIOWrapper.__init__(buffer) into a protocol

python · srittau · Feb 14, 2024 · Feb 14, 2024 · Feb 14, 2024 · Feb 14, 2024
commit 9ea524bd8863cca19accf1e3e5ae446054e7f981
diff --git a/stdlib/io.pyi b/stdlib/io.pyi
@@ -6,7 +6,7 @@ from _typeshed import FileDescriptorOrPath, ReadableBuffer, WriteableBuffer
 from collections.abc import Callable, Iterable, Iterator
 from os import _Opener
 from types import TracebackType
-from typing import IO, Any, BinaryIO, Literal, TextIO, TypeVar, overload
+from typing import IO, Any, BinaryIO, Literal, Protocol, TextIO, TypeVar, overload, type_check_only
 from typing_extensions import Self
 
 __all__ = [
@@ -146,10 +146,29 @@ class TextIOBase(IOBase):
     def readlines(self, __hint: int = -1) -> list[str]: ...  # type: ignore[override]
     def read(self, __size: int | None = ...) -> str: ...
 
+@type_check_only
+class _WrappedBuffer(Protocol):
+    name: str
+    closed: bool
+    def read(self, size: int = ...) -> bytes: ...
+    # Optional: def read1(size: int, /) -> bytes: ...
+    def write(self, b: bytes, /) -> object: ...
+    def flush(self) -> object: ...
+    def close(self) -> object: ...
+    def seekable() -> bool: ...
+    def readable() -> bool: ...
+    def writable() -> bool: ...
+    def truncate(size: int, /) -> int: ...
+    def fileno() -> int: ...
+    def isatty() -> int: ...
+    # Optional: Only needs to be present if seekable() returns True.
+    # def seek(self, offset: Literal[0], whence: Literal[2]) -> int: ...
+    # def tell(self) -> int: ...
+
 class TextIOWrapper(TextIOBase, TextIO):  # type: ignore[misc]  # incompatible definitions of write in the base classes
     def __init__(
         self,
-        buffer: IO[bytes],
+        buffer: _WrappedBuffer,
         encoding: str | None = ...,
         errors: str | None = ...,
         newline: str | None = ...,
@@ -180,7 +199,10 @@ class TextIOWrapper(TextIOBase, TextIO):  # type: ignore[misc]  # incompatible d
     def writelines(self, __lines: Iterable[str]) -> None: ...  # type: ignore[override]
     def readline(self, __size: int = -1) -> str: ...  # type: ignore[override]
     def readlines(self, __hint: int = -1) -> list[str]: ...  # type: ignore[override]
-    def seek(self, __cookie: int, __whence: int = 0) -> int: ...  # stubtest needs this
+    @overload
+    def seek(self, __cookie: int, __whence: Literal[0] = 0) -> int: ...
+    @overload
+    def seek(self, __cookie: Literal[0], __whence: Literal[1, 2]) -> int: ...
 
 class StringIO(TextIOWrapper):
     def __init__(self, initial_value: str | None = ..., newline: str | None = ...) -> None: ...