pyiter.transform

  1from abc import ABC, abstractmethod
  2from typing import Any, Generic, Iterable, Iterator, List, Optional, TypeVar, Generator
  3
  4
  5T = TypeVar("T")
  6
  7U = TypeVar("U")
  8
  9O = TypeVar("O")
 10
 11K = TypeVar("K")
 12
 13
 14class Transform(ABC, Generic[T, U], Iterable[U]):
 15    """A transform that applies a function to an iterable."""
 16
 17    iter: Iterable[T]
 18    cache: Optional[List[U]]
 19
 20    def __init__(self, iter: Iterable[T]):
 21        from .sequence import Sequence
 22
 23        self.iter = iter.__transform__ if isinstance(iter, Sequence) else iter
 24        self.cache = None
 25
 26    def __iter__(self) -> Iterator[U]:
 27        if self.cache:
 28            yield from self.cache
 29        else:
 30            cache: List[U] = []
 31            for x in self.__do_iter__():
 32                cache.append(x)
 33                yield x
 34            self.cache = cache
 35
 36    def __len__(self) -> int:
 37        # if not isinstance(self.iter, Transform):
 38        #     # not Sequence, just a wrapper of List, Tuple.etc.
 39        #     # we can get lenght of it directly.
 40        #     if hasattr(self.iter, '__len__'):
 41        #         return len(self.iter) # type: ignore
 42        #     elif hasattr(self.iter, '__length_hint__'):
 43        #         return self._iter.__length_hint__() # type: ignore
 44        # we need iterate all to get length
 45        cache = self.cache
 46        if cache is None:
 47            for _ in self:
 48                pass
 49            cache = self.cache
 50        if cache is not None:
 51            return len(cache)
 52        return 0
 53
 54    def __repr__(self) -> str:
 55        return self.__class__.__name__
 56
 57    @abstractmethod
 58    def __do_iter__(self) -> Iterator[U]:
 59        raise NotImplementedError
 60
 61    def transforms(self) -> "Iterable[Transform[Any, Any]]":
 62        inner = self.iter
 63        if isinstance(inner, Transform):
 64            yield from inner.transforms()
 65        yield self
 66
 67
 68class NonTransform(Transform[T, T]):
 69    """
 70    A [Transform] that does not transform the values.
 71    """
 72
 73    def __init__(self, iter: Iterable[T]) -> None:
 74        super().__init__(iter)
 75        self.cache = iter if isinstance(iter, list) else [*iter]
 76
 77    def __repr__(self) -> str:
 78        return f"{self.__class__.__name__}[{self.iter.__class__.__name__}]"
 79
 80    def __do_iter__(self) -> Iterator[T]:
 81        yield from self.iter
 82
 83    def __len__(self) -> int:
 84        if not isinstance(self.iter, Transform):
 85            # not Sequence, just a wrapper of List, Tuple.etc.
 86            # we can get lenght of it directly.
 87            if hasattr(self.iter, "__len__"):
 88                return len(self.iter)  # type: ignore
 89            elif hasattr(self.iter, "__length_hint__"):
 90                return self.iter.__length_hint__()  # type: ignore
 91        return super().__len__()
 92
 93
 94class InfinityTransform(Transform[T, T]):
 95    """Transform that iterates over an infinite sequence."""
 96
 97    def __init__(self, iter: Iterable[T]) -> None:
 98        super().__init__(iter)
 99
100    def __repr__(self) -> str:
101        return f"{self.__class__.__name__}[{self.iter.__class__.__name__}]"
102
103    def __do_iter__(self) -> Iterator[T]:
104        yield from self.iter
105
106    def __len__(self) -> int:
107        if self.cache is not None:
108            return len(self.cache)
109        raise OverflowError("Cannot determine the length of an infinite sequence.")
110
111
112def new_transform(iter: Iterable[T]) -> Transform[Any, T]:
113    from .sequence import Sequence
114
115    if isinstance(iter, Sequence):
116        return iter.__transform__
117    if isinstance(iter, Transform):
118        return iter  # type: ignore
119
120    if isinstance(iter, Generator):
121        return InfinityTransform(iter)  # type: ignore
122
123    return NonTransform(iter)