diff options
| -rw-r--r-- | utils.py | 46 |
1 files changed, 22 insertions, 24 deletions
@@ -1,55 +1,53 @@ +from __future__ import annotations from collections import deque -from collections.abc import Iterable, Iterator, Generator +from collections.abc import Iterable, Iterator from dataclasses import dataclass from itertools import islice from math import sqrt from sys import argv from typing import TypeVar, Generic, Union, Optional, cast -Numeric = Union[int, float] +NumT = TypeVar('NumT', float, int) -T = TypeVar('T', bound=Numeric) @dataclass(frozen=True) -class Point2D(Generic[T]): - x: T - y: T +class Point2D(Generic[NumT]): + x: NumT + y: NumT def __str__(self) -> str: return f'({self.x}, {self.y})' - def __add__(self, other: 'Point2D') -> Union['Point2D', bool]: + def __add__(self, other: Point2D[NumT]) -> Union[Point2D[NumT], bool]: if isinstance(other, Point2D): return Point2D(self.x + other.x, self.y + other.y) return NotImplemented - def __sub__(self, other: 'Point2D') -> Union['Point2D', bool]: + def __sub__(self, other: Point2D[NumT]) -> Union[Point2D[NumT], bool]: if isinstance(other, Point2D): return Point2D(self.x - other.x, self.y - other.y) return NotImplemented - def __abs__(self) -> 'Point2D': - return Point2D(abs(self.x), abs(self.y)) + def __abs__(self) -> Point2D[NumT]: + return Point2D(cast(NumT, abs(self.x)), cast(NumT, abs(self.y))) -T = TypeVar('T', bound=Numeric) -def norm_1(p: Point2D[T]) -> T: - return abs(p.x) + abs(p.y) +def norm_1(p: Point2D[NumT]) -> NumT: + return cast(NumT, abs(p.x) + abs(p.y)) def norm_2(p: Point2D) -> float: return sqrt(p.x ** 2 + p.y ** 2) -T = TypeVar('T', bound=Numeric) -def norm_inf(p: Point2D[T]) -> T: - return max(abs(p.x), abs(p.y)) +def norm_inf(p: Point2D[NumT]) -> NumT: + return cast(NumT, max(abs(p.x), abs(p.y))) -T = TypeVar('T', bound=Numeric) -def inbounds(p: Point2D[T], a: Point2D[T], b: Optional[Point2D[T]] = None) -> bool: - if b is None and isinstance(a.x, int): +def inbounds(p: Point2D[NumT], a: Point2D[NumT], _b: Optional[Point2D[NumT]] = None) -> bool: + if isinstance(_b, Point2D): + b = _b + else: b = a - a = cast(Point2D[T], Point2D(0, 0)) + a = cast(Point2D[NumT], Point2D(0, 0)) return p.x >= a.x and p.y >= a.y and p.x < b.x and p.y < b.y -def adjacent(p: Point2D[int], diagonal: bool = True) \ - -> Generator[Point2D[int], None, None]: +def adjacent(p: Point2D[int], diagonal: bool = True) -> Iterator[Point2D[int]]: for dx in range(-1, 2): for dy in range(-1, 2): if dx == 0 and dy == 0: continue @@ -57,11 +55,11 @@ def adjacent(p: Point2D[int], diagonal: bool = True) \ yield Point2D(p.x + dx, p.y + dy) def adjacent_bounded(p: Point2D[int], bound: Point2D[int], diagonal: bool = True) \ - -> filter[Point2D[int]]: # This has to return filter[...] not Generator[...] because ??? + -> Iterator[Point2D[int]]: return filter(lambda p: inbounds(p, bound), adjacent(p, diagonal)) T = TypeVar('T') -def sliding_window(iterable: Iterable[T], n: int) -> Generator[tuple[T, ...], None, None]: +def sliding_window(iterable: Iterable[T], n: int) -> Iterator[tuple[T, ...]]: # sliding_window('ABCDEFG', 4) -> ABCD BCDE CDEF DEFG it: Iterator[T] = iter(iterable) window: deque[T] = deque(islice(it, n), maxlen=n) |