Added type hints

This commit is contained in:
Andrew Murray 2023-12-24 15:59:12 +11:00
parent 41e45b5f95
commit d2dc8d3a6e

View File

@ -16,6 +16,10 @@
## ##
from __future__ import annotations from __future__ import annotations
from typing import Callable
from . import Image
class Iterator: class Iterator:
""" """
@ -29,14 +33,14 @@ class Iterator:
:param im: An image object. :param im: An image object.
""" """
def __init__(self, im): def __init__(self, im: Image.Image):
if not hasattr(im, "seek"): if not hasattr(im, "seek"):
msg = "im must have seek method" msg = "im must have seek method"
raise AttributeError(msg) raise AttributeError(msg)
self.im = im self.im = im
self.position = getattr(self.im, "_min_frame", 0) self.position = getattr(self.im, "_min_frame", 0)
def __getitem__(self, ix): def __getitem__(self, ix: int) -> Image.Image:
try: try:
self.im.seek(ix) self.im.seek(ix)
return self.im return self.im
@ -44,10 +48,10 @@ class Iterator:
msg = "end of sequence" msg = "end of sequence"
raise IndexError(msg) from e raise IndexError(msg) from e
def __iter__(self): def __iter__(self) -> Iterator:
return self return self
def __next__(self): def __next__(self) -> Image.Image:
try: try:
self.im.seek(self.position) self.im.seek(self.position)
self.position += 1 self.position += 1
@ -57,7 +61,10 @@ class Iterator:
raise StopIteration(msg) from e raise StopIteration(msg) from e
def all_frames(im, func=None): def all_frames(
im: Image.Image | list[Image.Image],
func: Callable[[Image.Image], Image.Image] | None = None,
) -> list[Image.Image]:
""" """
Applies a given function to all frames in an image or a list of images. Applies a given function to all frames in an image or a list of images.
The frames are returned as a list of separate images. The frames are returned as a list of separate images.