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 typing import Callable
from . import Image
class Iterator:
"""
@ -29,14 +33,14 @@ class Iterator:
:param im: An image object.
"""
def __init__(self, im):
def __init__(self, im: Image.Image):
if not hasattr(im, "seek"):
msg = "im must have seek method"
raise AttributeError(msg)
self.im = im
self.position = getattr(self.im, "_min_frame", 0)
def __getitem__(self, ix):
def __getitem__(self, ix: int) -> Image.Image:
try:
self.im.seek(ix)
return self.im
@ -44,10 +48,10 @@ class Iterator:
msg = "end of sequence"
raise IndexError(msg) from e
def __iter__(self):
def __iter__(self) -> Iterator:
return self
def __next__(self):
def __next__(self) -> Image.Image:
try:
self.im.seek(self.position)
self.position += 1
@ -57,7 +61,10 @@ class Iterator:
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.
The frames are returned as a list of separate images.