Merge pull request #7944 from nulano/type-image-open

radarhere · web-flow · commit 84a02c851aef · 2024-04-07T08:31:55.000+10:00
Add type hints for `Image.open`, `Image.init`, and `Image.Image.save`
diff --git a/src/PIL/BlpImagePlugin.py b/src/PIL/BlpImagePlugin.py
@@ -241,7 +241,7 @@ class BLPFormatError(NotImplementedError):
     pass
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] in (b"BLP1", b"BLP2")
 
 
diff --git a/src/PIL/BmpImagePlugin.py b/src/PIL/BmpImagePlugin.py
@@ -48,7 +48,7 @@
 }
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:2] == b"BM"
 
 
diff --git a/src/PIL/BufrStubImagePlugin.py b/src/PIL/BufrStubImagePlugin.py
@@ -29,7 +29,7 @@ def register_handler(handler):
 # Image adapter
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] == b"BUFR" or prefix[:4] == b"ZCZC"
 
 
diff --git a/src/PIL/CurImagePlugin.py b/src/PIL/CurImagePlugin.py
@@ -25,7 +25,7 @@
 # --------------------------------------------------------------------
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] == b"\0\0\2\0"
 
 
diff --git a/src/PIL/DcxImagePlugin.py b/src/PIL/DcxImagePlugin.py
@@ -29,7 +29,7 @@
 MAGIC = 0x3ADE68B1  # QUIZ: what's this value, then?
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return len(prefix) >= 4 and i32(prefix) == MAGIC
 
 
diff --git a/src/PIL/DdsImagePlugin.py b/src/PIL/DdsImagePlugin.py
@@ -562,7 +562,7 @@ def _save(im, fp, filename):
     )
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] == b"DDS "
 
 
diff --git a/src/PIL/EpsImagePlugin.py b/src/PIL/EpsImagePlugin.py
@@ -195,7 +195,7 @@ def readline(self):
         return b"".join(s).decode("latin-1")
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] == b"%!PS" or (len(prefix) >= 4 and i32(prefix) == 0xC6D3D0C5)
 
 
diff --git a/src/PIL/FliImagePlugin.py b/src/PIL/FliImagePlugin.py
@@ -27,7 +27,7 @@
 # decoder
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return (
         len(prefix) >= 6
         and i16(prefix, 4) in [0xAF11, 0xAF12]
diff --git a/src/PIL/FpxImagePlugin.py b/src/PIL/FpxImagePlugin.py
@@ -41,7 +41,7 @@
 # --------------------------------------------------------------------
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:8] == olefile.MAGIC
 
 
diff --git a/src/PIL/FtexImagePlugin.py b/src/PIL/FtexImagePlugin.py
@@ -107,7 +107,7 @@ def load_seek(self, pos):
         pass
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] == MAGIC
 
 
diff --git a/src/PIL/GbrImagePlugin.py b/src/PIL/GbrImagePlugin.py
@@ -29,7 +29,7 @@
 from ._binary import i32be as i32
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return len(prefix) >= 8 and i32(prefix, 0) >= 20 and i32(prefix, 4) in (1, 2)
 
 
diff --git a/src/PIL/GifImagePlugin.py b/src/PIL/GifImagePlugin.py
@@ -60,7 +60,7 @@ class LoadingStrategy(IntEnum):
 # Identify/read GIF files
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:6] in [b"GIF87a", b"GIF89a"]
 
 
diff --git a/src/PIL/GribStubImagePlugin.py b/src/PIL/GribStubImagePlugin.py
@@ -29,7 +29,7 @@ def register_handler(handler):
 # Image adapter
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] == b"GRIB" and prefix[7] == 1
 
 
diff --git a/src/PIL/Hdf5StubImagePlugin.py b/src/PIL/Hdf5StubImagePlugin.py
@@ -29,7 +29,7 @@ def register_handler(handler):
 # Image adapter
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:8] == b"\x89HDF\r\n\x1a\n"
 
 
diff --git a/src/PIL/IcnsImagePlugin.py b/src/PIL/IcnsImagePlugin.py
@@ -374,7 +374,7 @@ def _save(im, fp, filename):
         fp.flush()
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] == MAGIC
 
 
diff --git a/src/PIL/IcoImagePlugin.py b/src/PIL/IcoImagePlugin.py
@@ -114,7 +114,7 @@ def _save(im, fp, filename):
         fp.seek(current)
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] == _MAGIC
 
 
diff --git a/src/PIL/Image.py b/src/PIL/Image.py
@@ -41,7 +41,7 @@
 from collections.abc import Callable, MutableMapping
 from enum import IntEnum
 from types import ModuleType
-from typing import IO, TYPE_CHECKING, Any
+from typing import IO, TYPE_CHECKING, Any, Literal, cast
 
 # VERSION was removed in Pillow 6.0.0.
 # PILLOW_VERSION was removed in Pillow 9.0.0.
@@ -55,7 +55,7 @@
     _plugins,
 )
 from ._binary import i32le, o32be, o32le
-from ._typing import TypeGuard
+from ._typing import StrOrBytesPath, TypeGuard
 from ._util import DeferredError, is_path
 
 ElementTree: ModuleType | None
@@ -223,7 +223,7 @@ class Quantize(IntEnum):
     str,
     tuple[
         Callable[[IO[bytes], str | bytes], ImageFile.ImageFile],
-        Callable[[bytes], bool] | None,
+        Callable[[bytes], bool | str] | None,
     ],
 ] = {}
 MIME: dict[str, str] = {}
@@ -357,7 +357,7 @@ def preinit() -> None:
     _initialized = 1
 
 
-def init():
+def init() -> bool:
     """
     Explicitly initializes the Python Imaging Library. This function
     loads all available file format drivers.
@@ -368,7 +368,7 @@ def init():
 
     global _initialized
     if _initialized >= 2:
-        return 0
+        return False
 
     parent_name = __name__.rpartition(".")[0]
     for plugin in _plugins:
@@ -380,7 +380,8 @@ def init():
 
     if OPEN or SAVE:
         _initialized = 2
-        return 1
+        return True
+    return False
 
 
 # --------------------------------------------------------------------
@@ -2373,7 +2374,9 @@ def transform(x, y, matrix):
             (w, h), Transform.AFFINE, matrix, resample, fillcolor=fillcolor
         )
 
-    def save(self, fp, format=None, **params) -> None:
+    def save(
+        self, fp: StrOrBytesPath | IO[bytes], format: str | None = None, **params: Any
+    ) -> None:
         """
         Saves this image under the given filename.  If no format is
         specified, the format to use is determined from the filename
@@ -2454,6 +2457,8 @@ def save(self, fp, format=None, **params) -> None:
                 fp = builtins.open(filename, "r+b")
             else:
                 fp = builtins.open(filename, "w+b")
+        else:
+            fp = cast(IO[bytes], fp)
 
         try:
             save_handler(self, fp, filename)
@@ -3222,7 +3227,11 @@ def _decompression_bomb_check(size: tuple[int, int]) -> None:
         )
 
 
-def open(fp, mode="r", formats=None) -> Image:
+def open(
+    fp: StrOrBytesPath | IO[bytes],
+    mode: Literal["r"] = "r",
+    formats: list[str] | tuple[str, ...] | None = None,
+) -> ImageFile.ImageFile:
     """
     Opens and identifies the given image file.
 
@@ -3253,10 +3262,10 @@ def open(fp, mode="r", formats=None) -> Image:
     """
 
     if mode != "r":
-        msg = f"bad mode {repr(mode)}"
+        msg = f"bad mode {repr(mode)}"  # type: ignore[unreachable]
         raise ValueError(msg)
     elif isinstance(fp, io.StringIO):
-        msg = (
+        msg = (  # type: ignore[unreachable]
             "StringIO cannot be used to open an image. "
             "Binary data must be used instead."
         )
@@ -3265,7 +3274,7 @@ def open(fp, mode="r", formats=None) -> Image:
     if formats is None:
         formats = ID
     elif not isinstance(formats, (list, tuple)):
-        msg = "formats must be a list or tuple"
+        msg = "formats must be a list or tuple"  # type: ignore[unreachable]
         raise TypeError(msg)
 
     exclusive_fp = False
@@ -3276,6 +3285,8 @@ def open(fp, mode="r", formats=None) -> Image:
     if filename:
         fp = builtins.open(filename, "rb")
         exclusive_fp = True
+    else:
+        fp = cast(IO[bytes], fp)
 
     try:
         fp.seek(0)
@@ -3287,17 +3298,22 @@ def open(fp, mode="r", formats=None) -> Image:
 
     preinit()
 
-    accept_warnings = []
+    accept_warnings: list[str] = []
 
-    def _open_core(fp, filename, prefix, formats):
+    def _open_core(
+        fp: IO[bytes],
+        filename: str | bytes,
+        prefix: bytes,
+        formats: list[str] | tuple[str, ...],
+    ) -> ImageFile.ImageFile | None:
         for i in formats:
             i = i.upper()
             if i not in OPEN:
                 init()
             try:
                 factory, accept = OPEN[i]
                 result = not accept or accept(prefix)
-                if type(result) in [str, bytes]:
+                if isinstance(result, str):
                     accept_warnings.append(result)
                 elif result:
                     fp.seek(0)
@@ -3318,7 +3334,7 @@ def _open_core(fp, filename, prefix, formats):
     im = _open_core(fp, filename, prefix, formats)
 
     if im is None and formats is ID:
-        checked_formats = formats.copy()
+        checked_formats = ID.copy()
         if init():
             im = _open_core(
                 fp,
@@ -3448,7 +3464,7 @@ def merge(mode, bands):
 def register_open(
     id,
     factory: Callable[[IO[bytes], str | bytes], ImageFile.ImageFile],
-    accept: Callable[[bytes], bool] | None = None,
+    accept: Callable[[bytes], bool | str] | None = None,
 ) -> None:
     """
     Register an image file plugin.  This function should not be used
diff --git a/src/PIL/Jpeg2KImagePlugin.py b/src/PIL/Jpeg2KImagePlugin.py
@@ -317,7 +317,7 @@ def load(self):
         return ImageFile.ImageFile.load(self)
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return (
         prefix[:4] == b"\xff\x4f\xff\x51"
         or prefix[:12] == b"\x00\x00\x00\x0cjP  \x0d\x0a\x87\x0a"
diff --git a/src/PIL/JpegImagePlugin.py b/src/PIL/JpegImagePlugin.py
@@ -344,7 +344,7 @@ def DQT(self, marker):
 }
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     # Magic number was taken from https://en.wikipedia.org/wiki/JPEG
     return prefix[:3] == b"\xFF\xD8\xFF"
 
diff --git a/src/PIL/MicImagePlugin.py b/src/PIL/MicImagePlugin.py
@@ -25,7 +25,7 @@
 # --------------------------------------------------------------------
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:8] == olefile.MAGIC
 
 
diff --git a/src/PIL/PngImagePlugin.py b/src/PIL/PngImagePlugin.py
@@ -689,7 +689,7 @@ def chunk_fdAT(self, pos, length):
 # PNG reader
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:8] == _MAGIC
 
 
diff --git a/src/PIL/PsdImagePlugin.py b/src/PIL/PsdImagePlugin.py
@@ -44,7 +44,7 @@
 # read PSD images
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] == b"8BPS"
 
 
diff --git a/src/PIL/QoiImagePlugin.py b/src/PIL/QoiImagePlugin.py
@@ -13,7 +13,7 @@
 from ._binary import i32be as i32
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] == b"qoif"
 
 
diff --git a/src/PIL/TiffImagePlugin.py b/src/PIL/TiffImagePlugin.py
@@ -277,7 +277,7 @@
 ]
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return prefix[:4] in PREFIXES
 
 
diff --git a/src/PIL/WebPImagePlugin.py b/src/PIL/WebPImagePlugin.py
@@ -23,7 +23,7 @@
 }
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool | str:
     is_riff_file_format = prefix[:4] == b"RIFF"
     is_webp_file = prefix[8:12] == b"WEBP"
     is_valid_vp8_mode = prefix[12:16] in _VP8_MODES_BY_IDENTIFIER
@@ -34,6 +34,7 @@ def _accept(prefix):
                 "image file could not be identified because WEBP support not installed"
             )
         return True
+    return False
 
 
 class WebPImageFile(ImageFile.ImageFile):
diff --git a/src/PIL/WmfImagePlugin.py b/src/PIL/WmfImagePlugin.py
@@ -65,7 +65,7 @@ def load(self, im):
 # Read WMF file
 
 
-def _accept(prefix):
+def _accept(prefix: bytes) -> bool:
     return (
         prefix[:6] == b"\xd7\xcd\xc6\x9a\x00\x00" or prefix[:4] == b"\x01\x00\x00\x00"
     )
diff --git a/src/PIL/XpmImagePlugin.py b/src/PIL/XpmImagePlugin.py

Original file line number	Diff line number	Diff line change
`@@ -48,7 +48,7 @@`
`48`	`48`	`}`
`49`	`49`
`50`	`50`
`51`		`-def _accept(prefix):`
	`51`	`+def _accept(prefix: bytes) -> bool:`
`52`	`52`	`return prefix[:2] == b"BM"`
`53`	`53`
`54`	`54`
Original file line number	Diff line number	Diff line change
`@@ -562,7 +562,7 @@ def _save(im, fp, filename):`
`562`	`562`	`)`
`563`	`563`
`564`	`564`
`565`		`-def _accept(prefix):`
	`565`	`+def _accept(prefix: bytes) -> bool:`
`566`	`566`	`return prefix[:4] == b"DDS "`
`567`	`567`
`568`	`568`