From c85eb0cae5568ba9286ac15c3e8f1709f7bbe9e0 Mon Sep 17 00:00:00 2001 From: Andrew Murray Date: Sun, 28 Jul 2024 12:53:02 +1000 Subject: [PATCH] Added type hints --- Tests/test_file_png.py | 2 ++ src/PIL/Image.py | 30 +++++++++++++++++------------- src/PIL/ImageDraw2.py | 14 +++++++------- src/PIL/PngImagePlugin.py | 38 +++++++++++++++++++------------------- src/PIL/TiffImagePlugin.py | 8 +++++--- 5 files changed, 50 insertions(+), 42 deletions(-) diff --git a/Tests/test_file_png.py b/Tests/test_file_png.py index 4958b2222..0abf9866f 100644 --- a/Tests/test_file_png.py +++ b/Tests/test_file_png.py @@ -424,8 +424,10 @@ class TestFilePng: im = roundtrip(im, pnginfo=info) assert im.info == {"spam": "Eggs", "eggs": "Spam"} assert im.text == {"spam": "Eggs", "eggs": "Spam"} + assert isinstance(im.text["spam"], PngImagePlugin.iTXt) assert im.text["spam"].lang == "en" assert im.text["spam"].tkey == "Spam" + assert isinstance(im.text["eggs"], PngImagePlugin.iTXt) assert im.text["eggs"].lang == "en" assert im.text["eggs"].tkey == "Eggs" diff --git a/src/PIL/Image.py b/src/PIL/Image.py index 72d167a8a..eb3dd5a36 100644 --- a/src/PIL/Image.py +++ b/src/PIL/Image.py @@ -38,7 +38,7 @@ import struct import sys import tempfile import warnings -from collections.abc import Callable, MutableMapping, Sequence +from collections.abc import Callable, Iterator, MutableMapping, Sequence from enum import IntEnum from types import ModuleType from typing import ( @@ -744,11 +744,11 @@ class Image: new["shape"], new["typestr"] = _conv_type_shape(self) return new - def __getstate__(self): + def __getstate__(self) -> list[Any]: im_data = self.tobytes() # load image first return [self.info, self.mode, self.size, self.getpalette(), im_data] - def __setstate__(self, state) -> None: + def __setstate__(self, state: list[Any]) -> None: Image.__init__(self) info, mode, size, palette, data = state self.info = info @@ -1683,7 +1683,9 @@ class Image: x, y = self.im.getprojection() return list(x), list(y) - def histogram(self, mask: Image | None = None, extrema=None) -> list[int]: + def histogram( + self, mask: Image | None = None, extrema: tuple[float, float] | None = None + ) -> list[int]: """ Returns a histogram for the image. The histogram is returned as a list of pixel counts, one for each pixel value in the source @@ -1709,12 +1711,14 @@ class Image: mask.load() return self.im.histogram((0, 0), mask.im) if self.mode in ("I", "F"): - if extrema is None: - extrema = self.getextrema() - return self.im.histogram(extrema) + return self.im.histogram( + extrema if extrema is not None else self.getextrema() + ) return self.im.histogram() - def entropy(self, mask: Image | None = None, extrema=None): + def entropy( + self, mask: Image | None = None, extrema: tuple[float, float] | None = None + ) -> float: """ Calculates and returns the entropy for the image. @@ -1735,9 +1739,9 @@ class Image: mask.load() return self.im.entropy((0, 0), mask.im) if self.mode in ("I", "F"): - if extrema is None: - extrema = self.getextrema() - return self.im.entropy(extrema) + return self.im.entropy( + extrema if extrema is not None else self.getextrema() + ) return self.im.entropy() def paste( @@ -3881,7 +3885,7 @@ class Exif(_ExifBase): # returns a dict with any single item tuples/lists as individual values return {k: self._fixup(v) for k, v in src_dict.items()} - def _get_ifd_dict(self, offset: int, group=None): + def _get_ifd_dict(self, offset: int, group: int | None = None): try: # an offset pointer to the location of the nested embedded IFD. # It should be a long, but may be corrupted. @@ -4136,7 +4140,7 @@ class Exif(_ExifBase): else: del self._data[tag] - def __iter__(self): + def __iter__(self) -> Iterator[int]: keys = set(self._data) if self._info is not None: keys.update(self._info) diff --git a/src/PIL/ImageDraw2.py b/src/PIL/ImageDraw2.py index c0d89e4aa..1ad239382 100644 --- a/src/PIL/ImageDraw2.py +++ b/src/PIL/ImageDraw2.py @@ -24,7 +24,7 @@ """ from __future__ import annotations -from typing import AnyStr, BinaryIO +from typing import Any, AnyStr, BinaryIO from . import Image, ImageColor, ImageDraw, ImageFont, ImagePath from ._typing import Coords, StrOrBytesPath @@ -111,7 +111,7 @@ class Draw: (xoffset, yoffset) = offset self.transform = (1, 0, xoffset, 0, 1, yoffset) - def arc(self, xy: Coords, start, end, *options) -> None: + def arc(self, xy: Coords, start, end, *options: Any) -> None: """ Draws an arc (a portion of a circle outline) between the start and end angles, inside the given bounding box. @@ -120,7 +120,7 @@ class Draw: """ self.render("arc", xy, start, end, *options) - def chord(self, xy: Coords, start, end, *options) -> None: + def chord(self, xy: Coords, start, end, *options: Any) -> None: """ Same as :py:meth:`~PIL.ImageDraw2.Draw.arc`, but connects the end points with a straight line. @@ -129,7 +129,7 @@ class Draw: """ self.render("chord", xy, start, end, *options) - def ellipse(self, xy: Coords, *options) -> None: + def ellipse(self, xy: Coords, *options: Any) -> None: """ Draws an ellipse inside the given bounding box. @@ -137,7 +137,7 @@ class Draw: """ self.render("ellipse", xy, *options) - def line(self, xy: Coords, *options) -> None: + def line(self, xy: Coords, *options: Any) -> None: """ Draws a line between the coordinates in the ``xy`` list. @@ -145,7 +145,7 @@ class Draw: """ self.render("line", xy, *options) - def pieslice(self, xy: Coords, start, end, *options) -> None: + def pieslice(self, xy: Coords, start, end, *options: Any) -> None: """ Same as arc, but also draws straight lines between the end points and the center of the bounding box. @@ -154,7 +154,7 @@ class Draw: """ self.render("pieslice", xy, start, end, *options) - def polygon(self, xy: Coords, *options) -> None: + def polygon(self, xy: Coords, *options: Any) -> None: """ Draws a polygon. diff --git a/src/PIL/PngImagePlugin.py b/src/PIL/PngImagePlugin.py index 58db7777c..52cfcd13a 100644 --- a/src/PIL/PngImagePlugin.py +++ b/src/PIL/PngImagePlugin.py @@ -38,6 +38,7 @@ import re import struct import warnings import zlib +from collections.abc import Callable from enum import IntEnum from typing import IO, TYPE_CHECKING, Any, NamedTuple, NoReturn @@ -135,7 +136,7 @@ class Blend(IntEnum): """ -def _safe_zlib_decompress(s): +def _safe_zlib_decompress(s: bytes) -> bytes: dobj = zlib.decompressobj() plaintext = dobj.decompress(s, MAX_TEXT_CHUNK) if dobj.unconsumed_tail: @@ -783,7 +784,7 @@ class PngImageFile(ImageFile.ImageFile): self._mode = self.png.im_mode self._size = self.png.im_size self.info = self.png.im_info - self._text = None + self._text: dict[str, str | iTXt] | None = None self.tile = self.png.im_tile self.custom_mimetype = self.png.im_custom_mimetype self.n_frames = self.png.im_n_frames or 1 @@ -810,7 +811,7 @@ class PngImageFile(ImageFile.ImageFile): self.is_animated = self.n_frames > 1 @property - def text(self): + def text(self) -> dict[str, str | iTXt]: # experimental if self._text is None: # iTxt, tEXt and zTXt chunks may appear at the end of the file @@ -822,6 +823,7 @@ class PngImageFile(ImageFile.ImageFile): self.load() if self.is_animated: self.seek(frame) + assert self._text is not None return self._text def verify(self) -> None: @@ -1105,7 +1107,7 @@ def putchunk(fp: IO[bytes], cid: bytes, *data: bytes) -> None: class _idat: # wrap output from the encoder in IDAT chunks - def __init__(self, fp, chunk) -> None: + def __init__(self, fp: IO[bytes], chunk: Callable[..., None]) -> None: self.fp = fp self.chunk = chunk @@ -1116,7 +1118,7 @@ class _idat: class _fdat: # wrap encoder output in fdAT chunks - def __init__(self, fp: IO[bytes], chunk, seq_num: int) -> None: + def __init__(self, fp: IO[bytes], chunk: Callable[..., None], seq_num: int) -> None: self.fp = fp self.chunk = chunk self.seq_num = seq_num @@ -1135,7 +1137,7 @@ class _Frame(NamedTuple): def _write_multiple_frames( im: Image.Image, fp: IO[bytes], - chunk, + chunk: Callable[..., None], mode: str, rawmode: str, default_image: Image.Image | None, @@ -1275,7 +1277,11 @@ def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: def _save( - im: Image.Image, fp, filename: str | bytes, chunk=putchunk, save_all: bool = False + im: Image.Image, + fp: IO[bytes], + filename: str | bytes, + chunk: Callable[..., None] = putchunk, + save_all: bool = False, ) -> None: # save an image to disk (called by the save method) @@ -1483,22 +1489,16 @@ def _save( def getchunks(im: Image.Image, **params: Any) -> list[tuple[bytes, bytes, bytes]]: """Return a list of PNG chunks representing this image.""" + from io import BytesIO - class collector: - data = [] + chunks = [] - def write(self, data: bytes) -> None: - pass - - def append(self, chunk: tuple[bytes, bytes, bytes]) -> None: - self.data.append(chunk) - - def append(fp: collector, cid: bytes, *data: bytes) -> None: + def append(fp: IO[bytes], cid: bytes, *data: bytes) -> None: byte_data = b"".join(data) crc = o32(_crc32(byte_data, _crc32(cid))) - fp.append((cid, byte_data, crc)) + chunks.append((cid, byte_data, crc)) - fp = collector() + fp = BytesIO() try: im.encoderinfo = params @@ -1506,7 +1506,7 @@ def getchunks(im: Image.Image, **params: Any) -> list[tuple[bytes, bytes, bytes] finally: del im.encoderinfo - return fp.data + return chunks # -------------------------------------------------------------------- diff --git a/src/PIL/TiffImagePlugin.py b/src/PIL/TiffImagePlugin.py index ad050ffe4..318e9825c 100644 --- a/src/PIL/TiffImagePlugin.py +++ b/src/PIL/TiffImagePlugin.py @@ -1020,7 +1020,7 @@ class ImageFileDirectory_v1(ImageFileDirectory_v2): .. deprecated:: 3.0.0 """ - def __init__(self, *args, **kwargs) -> None: + def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._legacy_api = True @@ -1142,13 +1142,15 @@ class TiffImageFile(ImageFile.ImageFile): self._seek(0) @property - def n_frames(self): - if self._n_frames is None: + def n_frames(self) -> int: + current_n_frames = self._n_frames + if current_n_frames is None: current = self.tell() self._seek(len(self._frame_pos)) while self._n_frames is None: self._seek(self.tell() + 1) self.seek(current) + assert self._n_frames is not None return self._n_frames def seek(self, frame: int) -> None: