lazylist.py

from __future__ import annotations

import itertools
import warnings
from typing import (
    Any,
    Callable,
    Generator,
    Generic,
    Iterable,
    Iterator,
    List,
    Optional,
    Protocol,
    Tuple,
    TypeVar,
    Union,
    overload,
)

_a = TypeVar("_a")
_b = TypeVar("_b")


class SupportsDunderLT(Protocol):
    def __lt__(self, __other: Any) -> bool:
        ...


class SupportsDunderGT(Protocol):
    def __gt__(self, __other: Any) -> bool:
        ...


SupportsRichComparison = Union[SupportsDunderLT, SupportsDunderGT]
SupportsRichComparisonT = TypeVar(
    "SupportsRichComparisonT", bound=SupportsRichComparison
)


class ValueWarning(Warning):
    pass


class LazyList(Generic[_a]):
    def __init__(self, _iter: Iterable[_a]) -> None:
        self._iter = iter(_iter)
        self._list: List[_a] = []

    def __iter__(self) -> Iterator[_a]:
        yield from self._list
        yield from self._remaining_iter()

    def _remaining_iter(self) -> Iterator[_a]:
        for value in self._iter:
            self._list.append(value)
            yield value

    def __len__(self) -> int:
        return len(self._list) + sum(1 for _ in self._remaining_iter())

    @overload
    def __getitem__(self, idx: int) -> _a:
        ...

    @overload
    def __getitem__(self, idx: slice) -> LazyList[_a]:
        ...

    def __getitem__(self, idx: Union[int, slice]) -> Union[LazyList[_a], _a]:
        if isinstance(idx, int):
            if idx > (len(self._list) - 1):
                for value in self._remaining_iter():
                    if idx == (len(self._list) - 1):
                        return value
                else:
                    raise IndexError(
                        f"Sequence index {idx} out of range for list of length {len(self._list)}"
                    )
            return self._list[idx]

        return LazyList(itertools.islice(self, idx.start, idx.stop, idx.step))

    def __repr__(self) -> str:
        return f"{self.__class__.__qualname__}(length >= {len(self._list)})"

    def insert(self, idx: int, value: _a) -> LazyList[_a]:
        def _inserter() -> Generator[_a, None, None]:
            yield from self[:idx]
            yield value
            yield from self[idx:]

        return LazyList(_inserter())

    def extend(self, new_values: Iterable[_a]) -> LazyList[_a]:
        def _extender() -> Generator[_a, None, None]:
            yield from self
            yield from new_values

        return LazyList(_extender())

    def append(self, new_value: _a) -> LazyList[_a]:
        def _appender() -> Generator[_a, None, None]:
            yield from self
            yield new_value

        return LazyList(_appender())

    def remove(self, value: _a) -> LazyList[_a]:
        def _remover() -> Generator[_a, None, None]:
            found = False
            for _value in self:
                if not found and _value == value:
                    found = True
                    continue
                yield _value
            if not found:
                warnings.warn(
                    f"Value {value:r} not found in list, so not removed", ValueWarning
                )

        return LazyList(_remover())

    def pop(self, index: int) -> Tuple[_a, LazyList[_a]]:
        head = self[:index]
        popped = self[index]
        tail = self[index + 1 :]
        return popped, LazyList(head).extend(tail)

    def map(self, fun: Callable[[_a], _b]) -> LazyList[_b]:
        return LazyList(fun(val) for val in self)

    def reverse(self) -> LazyList[_a]:
        return LazyList(reversed(list(self)))

    def copy(self) -> LazyList[_a]:
        return LazyList(self)

    def count(self, value: _a) -> int:
        return sum(1 for _value in self if _value == value)

    def index(self, value: _a) -> int:
        for i, _value in enumerate(self):
            if _value == value:
                return i
        else:
            raise ValueError("Value not found")

    @overload
    def sort(
        self: LazyList[SupportsRichComparisonT],
        *,
        key: None = None,
        reverse: bool = False,
    ) -> LazyList[_a]:
        ...

    @overload
    def sort(
        self, *, key: Callable[[_a], SupportsRichComparison], reverse: bool = False
    ) -> LazyList[_a]:
        ...

    def sort(
        self: Union[LazyList[_a], LazyList[SupportsRichComparisonT]],
        *,
        key: Optional[Callable[[_a], SupportsRichComparison]] = None,
        reverse: bool = False,
    ) -> LazyList[_a]:
        return LazyList(sorted(self, key=key, reverse=reverse)) # type: ignore