diff --git a/src/affine/__init__.py b/src/affine/__init__.py index 41c54ed..ca7fe02 100644 --- a/src/affine/__init__.py +++ b/src/affine/__init__.py @@ -37,6 +37,7 @@ from collections.abc import MutableSequence, Sequence from functools import cached_property import math +from typing import overload import warnings from attrs import astuple, define, field @@ -550,6 +551,17 @@ def __add__(self, other): __iadd__ = __add__ + @overload + def __matmul__(self, other: Affine) -> Affine: ... + @overload + def __matmul__(self, other: tuple[float, float]) -> tuple[float, float]: ... + @overload + def __matmul__( + self, other: tuple[float, float, float] + ) -> tuple[float, float, float]: ... + # For other float sequences, we don't know the returned tuple length here + @overload + def __matmul__(self, other: Sequence[float]) -> tuple[float, ...]: ... def __matmul__(self, other): """Matrix multiplication. @@ -605,11 +617,15 @@ def __matmul__(self, other): def __rmatmul__(self, other): return NotImplemented - def __imatmul__(self, other): + def __imatmul__(self, other): # type: ignore if not isinstance(other, Affine): raise TypeError("Operation not supported") return NotImplemented + @overload + def __mul__(self, other: Affine) -> Affine: ... + @overload + def __mul__(self, other: tuple[float, float]) -> tuple[float, float]: ... def __mul__(self, other): """Multiplication. @@ -642,7 +658,7 @@ def __mul__(self, other): def __rmul__(self, other): return NotImplemented - def __imul__(self, other): + def __imul__(self, other): # type: ignore if isinstance(other, tuple): warnings.warn( "in-place multiplication with tuple is deprecated",