TemporaryFile(): Handle one more case correctly (python#10284)

sobolevn · Jun 14, 2023 · 5f9d05c · 5f9d05c
1 parent 6e90fce
commit 5f9d05c
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 0 deletions.
diff --git a/stdlib/tempfile.pyi b/stdlib/tempfile.pyi
@@ -191,6 +191,17 @@ else:
             errors: str | None = None,
         ) -> io.FileIO: ...
         @overload
+        def TemporaryFile(
+            *,
+            buffering: Literal[0],
+            encoding: str | None = None,
+            newline: str | None = None,
+            suffix: AnyStr | None = None,
+            prefix: AnyStr | None = None,
+            dir: GenericPath[AnyStr] | None = None,
+            errors: str | None = None,
+        ) -> io.FileIO: ...
+        @overload
         def TemporaryFile(
             mode: OpenBinaryModeWriting,
             buffering: Literal[-1, 1] = -1,
@@ -260,6 +271,16 @@ else:
             dir: GenericPath[AnyStr] | None = None,
         ) -> io.FileIO: ...
         @overload
+        def TemporaryFile(
+            *,
+            buffering: Literal[0],
+            encoding: str | None = None,
+            newline: str | None = None,
+            suffix: AnyStr | None = None,
+            prefix: AnyStr | None = None,
+            dir: GenericPath[AnyStr] | None = None,
+        ) -> io.FileIO: ...
+        @overload
         def TemporaryFile(
             mode: OpenBinaryModeUpdating = "w+b",
             buffering: Literal[-1, 1] = -1,

diff --git a/test_cases/stdlib/check_tempfile.py b/test_cases/stdlib/check_tempfile.py
@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import io
+import sys
+from tempfile import TemporaryFile, _TemporaryFileWrapper
+from typing_extensions import assert_type
+
+if sys.platform == "win32":
+    assert_type(TemporaryFile(), _TemporaryFileWrapper[bytes])
+    assert_type(TemporaryFile("w+"), _TemporaryFileWrapper[str])
+    assert_type(TemporaryFile("w+b"), _TemporaryFileWrapper[bytes])
+    assert_type(TemporaryFile("wb"), _TemporaryFileWrapper[bytes])
+    assert_type(TemporaryFile("rb"), _TemporaryFileWrapper[bytes])
+    assert_type(TemporaryFile("wb", 0), _TemporaryFileWrapper[bytes])
+    assert_type(TemporaryFile(mode="w+"), _TemporaryFileWrapper[str])
+    assert_type(TemporaryFile(mode="w+b"), _TemporaryFileWrapper[bytes])
+    assert_type(TemporaryFile(mode="wb"), _TemporaryFileWrapper[bytes])
+    assert_type(TemporaryFile(mode="rb"), _TemporaryFileWrapper[bytes])
+    assert_type(TemporaryFile(buffering=0), _TemporaryFileWrapper[bytes])
+else:
+    assert_type(TemporaryFile(), io.BufferedRandom)
+    assert_type(TemporaryFile("w+"), io.TextIOWrapper)
+    assert_type(TemporaryFile("w+b"), io.BufferedRandom)
+    assert_type(TemporaryFile("wb"), io.BufferedWriter)
+    assert_type(TemporaryFile("rb"), io.BufferedReader)
+    assert_type(TemporaryFile("wb", 0), io.FileIO)
+    assert_type(TemporaryFile(mode="w+"), io.TextIOWrapper)
+    assert_type(TemporaryFile(mode="w+b"), io.BufferedRandom)
+    assert_type(TemporaryFile(mode="wb"), io.BufferedWriter)
+    assert_type(TemporaryFile(mode="rb"), io.BufferedReader)
+    assert_type(TemporaryFile(buffering=0), io.FileIO)