__iter__ type annotation fixed for vector class

This commit is contained in:
Nicolas Kruse 2025-11-27 17:20:46 +01:00
parent 89e8efb864
commit 054ce6d507
1 changed files with 12 additions and 6 deletions

View File

@ -1,5 +1,5 @@
from . import variable from . import variable
from typing import Generic, TypeVar, Iterable, Any, overload, TypeAlias, Callable from typing import Generic, TypeVar, Iterable, Any, overload, TypeAlias, Callable, Iterator
import copapy as cp import copapy as cp
VecNumLike: TypeAlias = 'vector[int] | vector[float] | variable[int] | variable[float] | int | float | bool' VecNumLike: TypeAlias = 'vector[int] | vector[float] | variable[int] | variable[float] | int | float | bool'
@ -31,6 +31,12 @@ class vector(Generic[T]):
def __getitem__(self, index: int) -> variable[T] | T: def __getitem__(self, index: int) -> variable[T] | T:
return self.values[index] return self.values[index]
def __neg__(self) -> 'vector[float] | vector[int]':
return vector(-a for a in self.values)
def __iter__(self) -> Iterator[variable[T] | T]:
return iter(self.values)
@overload @overload
def __add__(self: 'vector[int]', other: VecFloatLike) -> 'vector[float]': ... def __add__(self: 'vector[int]', other: VecFloatLike) -> 'vector[float]': ...
@overload @overload
@ -163,11 +169,11 @@ class vector(Generic[T]):
mag = self.magnitude() + epsilon mag = self.magnitude() + epsilon
return self / mag return self / mag
def __neg__(self) -> 'vector[float] | vector[int]': def homogenize(self) -> 'vector[T]':
return vector(-a for a in self.values) if any(isinstance(val, variable) for val in self.values):
return vector(variable(val) if not isinstance(val, variable) else val for val in self.values)
def __iter__(self) -> Iterable[variable[T] | T]: else:
return iter(self.values) return self
def map(self, func: Callable[[Any], variable[U] | U]) -> 'vector[U]': def map(self, func: Callable[[Any], variable[U] | U]) -> 'vector[U]':
"""Applies a function to each element of the vector and returns a new vector.""" """Applies a function to each element of the vector and returns a new vector."""