Add a return type for PIL.Image.load (#9466)

Co-authored-by: Alex Waygood <[email protected]>

Author: DMRobertson

stubs/Pillow/PIL/EpsImagePlugin.pyi

@@ -1,6 +1,7 @@
 from typing import Any, ClassVar
 from typing_extensions import Literal
 
+from ._imaging import _PixelAccessor
 from .ImageFile import ImageFile
 
 split: Any
@@ -24,5 +25,5 @@ class EpsImageFile(ImageFile):
     im: Any
     mode: Any
     tile: Any
-    def load(self, scale: int = 1, transparency: bool = False) -> None: ...
+    def load(self, scale: int = 1, transparency: bool = False) -> _PixelAccessor: ...
     def load_seek(self, *args, **kwargs) -> None: ...

stubs/Pillow/PIL/FpxImagePlugin.pyi

@@ -2,6 +2,7 @@ from _typeshed import Incomplete
 from typing import Any, ClassVar
 from typing_extensions import Literal, TypeAlias
 
+from ._imaging import _PixelAccessor
 from .ImageFile import ImageFile
 
 _OleFileIO: TypeAlias = Any  # olefile.OleFileIO
@@ -19,4 +20,4 @@ class FpxImageFile(ImageFile):
     jpeg: dict[int, Incomplete]
     tile_prefix: Incomplete
     stream: list[str]
-    def load(self): ...
+    def load(self) -> _PixelAccessor: ...

stubs/Pillow/PIL/GbrImagePlugin.pyi

@@ -1,10 +1,11 @@
 from typing import Any, ClassVar
 from typing_extensions import Literal
 
+from ._imaging import _PixelAccessor
 from .ImageFile import ImageFile
 
 class GbrImageFile(ImageFile):
     format: ClassVar[Literal["GBR"]]
     format_description: ClassVar[str]
     im: Any
-    def load(self) -> None: ...
+    def load(self) -> _PixelAccessor: ...

stubs/Pillow/PIL/IcnsImagePlugin.pyi

@@ -2,6 +2,7 @@ from _typeshed import Incomplete
 from typing import Any, ClassVar
 from typing_extensions import Literal
 
+from ._imaging import _PixelAccessor
 from .ImageFile import ImageFile
 
 enable_jpeg2k: Any
@@ -33,4 +34,4 @@ class IcnsImageFile(ImageFile):
     best_size: Any
     im: Any
     mode: Any
-    def load(self) -> None: ...
+    def load(self) -> _PixelAccessor: ...

stubs/Pillow/PIL/IcoImagePlugin.pyi

@@ -1,6 +1,7 @@
 from typing import Any, ClassVar
 from typing_extensions import Literal
 
+from ._imaging import _PixelAccessor
 from .ImageFile import ImageFile
 
 class IcoFile:
@@ -22,5 +23,5 @@ class IcoImageFile(ImageFile):
     def size(self, value) -> None: ...
     im: Any
     mode: Any
-    def load(self) -> None: ...
+    def load(self) -> _PixelAccessor: ...
     def load_seek(self) -> None: ...

stubs/Pillow/PIL/Image.pyi

@@ -13,6 +13,7 @@ from ._imaging import (
     FIXED as FIXED,
     HUFFMAN_ONLY as HUFFMAN_ONLY,
     RLE as RLE,
+    _PixelAccessor,
 )
 from .ImageFilter import Filter
 from .ImagePalette import ImagePalette
@@ -180,7 +181,7 @@ class Image:
     def tobytes(self, encoder_name: str = "raw", *args) -> bytes: ...
     def tobitmap(self, name: str = "image") -> bytes: ...
     def frombytes(self, data: bytes, decoder_name: str = "raw", *args) -> None: ...
-    def load(self) -> None: ...
+    def load(self) -> _PixelAccessor: ...
     def verify(self) -> None: ...
     def convert(
         self,

stubs/Pillow/PIL/ImageFile.pyi

@@ -2,6 +2,7 @@ from _typeshed import Incomplete, Unused
 from typing import Any, NoReturn
 from typing_extensions import Self
 
+from ._imaging import _PixelAccessor
 from .Image import Image
 
 MAXBLOCK: int
@@ -24,12 +25,12 @@ class ImageFile(Image):
     def verify(self) -> None: ...
     map: Any
     im: Any
-    def load(self): ...
+    def load(self) -> _PixelAccessor: ...
     def load_prepare(self) -> None: ...
     def load_end(self) -> None: ...
 
 class StubImageFile(ImageFile):
-    def load(self) -> None: ...
+    def load(self) -> _PixelAccessor: ...
 
 class Parser:
     incremental: Incomplete | None

stubs/Pillow/PIL/IptcImagePlugin.pyi

@@ -1,6 +1,7 @@
 from typing import Any, ClassVar
 from typing_extensions import Literal
 
+from ._imaging import _PixelAccessor
 from .ImageFile import ImageFile
 
 COMPRESSION: Any
@@ -15,6 +16,6 @@ class IptcImageFile(ImageFile):
     def getint(self, key): ...
     def field(self): ...
     im: Any
-    def load(self): ...
+    def load(self) -> _PixelAccessor: ...
 
 def getiptcinfo(im): ...

stubs/Pillow/PIL/Jpeg2KImagePlugin.pyi

@@ -1,11 +1,12 @@
 from typing import Any, ClassVar
 from typing_extensions import Literal
 
+from ._imaging import _PixelAccessor
 from .ImageFile import ImageFile
 
 class Jpeg2KImageFile(ImageFile):
     format: ClassVar[Literal["JPEG2000"]]
     format_description: ClassVar[str]
     reduce: Any
     tile: Any
-    def load(self): ...
+    def load(self) -> _PixelAccessor: ...

stubs/Pillow/PIL/PngImagePlugin.pyi

@@ -102,7 +102,6 @@ class PngImageFile(ImageFile):
     def load_read(self, read_bytes): ...
     png: Any
     im: Any
-    pyaccess: Any
     def load_end(self) -> None: ...
     def getexif(self): ...
 

stubs/Pillow/PIL/PyAccess.pyi

@@ -1,20 +1,22 @@
 from logging import Logger
 from typing import Any
 
+from PIL._imaging import _PixelAccessor
+
 ffi: Any
 logger: Logger
 
-class PyAccess:
+class PyAccess(_PixelAccessor):
     readonly: Any
     image8: Any
     image32: Any
     image: Any
     def __init__(self, img, readonly: bool = False) -> None: ...
-    def __setitem__(self, xy, color) -> None: ...
-    def __getitem__(self, xy): ...
-    putpixel: Any
-    getpixel: Any
-    def check_xy(self, xy): ...
+    def __setitem__(self, xy: tuple[int, int], color) -> None: ...
+    def __getitem__(self, xy: tuple[int, int]) -> Any: ...
+    def putpixel(self, xy: tuple[int, int], color) -> None: ...
+    def getpixel(self, xy: tuple[int, int]) -> Any: ...
+    def check_xy(self, xy: tuple[int, int]): ...
 
 class _PyAccess32_2(PyAccess):
     def get_pixel(self, x, y): ...

stubs/Pillow/PIL/TiffImagePlugin.pyi

@@ -5,6 +5,7 @@ from types import TracebackType
 from typing import Any, ClassVar
 from typing_extensions import Literal
 
+from ._imaging import _PixelAccessor
 from .ImageFile import ImageFile
 
 logger: Any
@@ -153,7 +154,7 @@ class TiffImageFile(ImageFile):
     im: Any
     def seek(self, frame) -> None: ...
     def tell(self): ...
-    def load(self): ...
+    def load(self) -> _PixelAccessor: ...
     def load_end(self) -> None: ...
 
 SAVE_INFO: Any

stubs/Pillow/PIL/WalImageFile.pyi

@@ -2,11 +2,12 @@ from typing import ClassVar
 from typing_extensions import Literal
 
 from . import ImageFile
+from ._imaging import _PixelAccessor
 
 class WalImageFile(ImageFile.ImageFile):
     format: ClassVar[Literal["WAL"]]
     format_description: ClassVar[str]
-    def load(self) -> None: ...
+    def load(self) -> _PixelAccessor: ...
 
 def open(filename): ...
 

stubs/Pillow/PIL/WebPImagePlugin.pyi

@@ -1,6 +1,7 @@
 from typing import Any, ClassVar
 from typing_extensions import Literal, TypeAlias
 
+from ._imaging import _PixelAccessor
 from .ImageFile import ImageFile
 
 SUPPORTED: bool
@@ -13,5 +14,5 @@ class WebPImageFile(ImageFile):
     def seek(self, frame) -> None: ...
     fp: Any
     tile: Any
-    def load(self): ...
+    def load(self) -> _PixelAccessor: ...
     def tell(self): ...

stubs/Pillow/PIL/WmfImagePlugin.pyi

@@ -3,6 +3,7 @@ from _typeshed import Incomplete
 from typing import Any, ClassVar
 from typing_extensions import Literal
 
+from ._imaging import _PixelAccessor
 from .ImageFile import StubImageFile
 
 def register_handler(handler) -> None: ...
@@ -16,4 +17,4 @@ if sys.platform == "win32":
 class WmfStubImageFile(StubImageFile):
     format: ClassVar[Literal["WMF"]]
     format_description: ClassVar[str]
-    def load(self, dpi: Incomplete | None = None) -> None: ...
+    def load(self, dpi: Incomplete | None = None) -> _PixelAccessor: ...

stubs/Pillow/PIL/_imaging.pyi

@@ -1,5 +1,6 @@
 from _typeshed import Incomplete
 from collections.abc import Sequence
+from typing import Protocol
 from typing_extensions import Literal
 
 DEFAULT_STRATEGY: Literal[0]
@@ -8,6 +9,20 @@ HUFFMAN_ONLY: Literal[2]
 RLE: Literal[3]
 FIXED: Literal[4]
 
+class _PixelAccessor(Protocol):  # noqa: Y046
+    # PIL has two concrete types for accessing an image's pixels by coordinate lookup:
+    # PixelAccess (written in C; not runtime-importable) and PyAccess (written in
+    # Python + cffi; is runtime-importable). PixelAccess came first. PyAccess was added
+    # in later to support PyPy, but otherwise is intended to expose the same interface
+    # PixelAccess.
+    #
+    # This protocol describes that interface.
+    # TODO: should the color args and getter return types be _Color?
+    def __setitem__(self, xy: tuple[int, int], color: Incomplete) -> None: ...
+    def __getitem__(self, xy: tuple[int, int]) -> Incomplete: ...
+    def putpixel(self, xy: tuple[int, int], color: Incomplete) -> None: ...
+    def getpixel(self, xy: tuple[int, int]) -> Incomplete: ...
+
 class _Path:
     def __getattr__(self, item: str) -> Incomplete: ...