def test_peek(fiter, lst): initial_len = len_or_none(fiter) assert fiter.peek(5) == lst[:5] assert initial_len == len_or_none(fiter) assert fiter.peek() == lst[0] assert_equal_it(fiter, lst) assert fiter.peek(0) == []
def test_combinations_with_replacement(fiter, lst): init_len = len_or_none(fiter) combs = fiter.combinations_with_replacement(3) if init_len is not None: assert len_or_none(combs) is not None assert_equal_it(combs, combinations_with_replacement(lst, 3))
def fiter(request, lst): if request.param == "list": fiter = FIt(lst) assert len_or_none(fiter) is not None return fiter else: fiter = FIt(iter(lst)) assert len_or_none(fiter) is None return fiter
def tee(self, n: int = 2): """See itertools.tee_ .. _itertools.tee: https://docs.python.org/3/library/itertools.html#itertools.tee """ # noqa length = len_or_none(self) return [FIt(t, length) for t in tee(self, n)]
def zip(self, *iterables, longest=False, fill_value=None, strict=False) -> FIt[Tuple]: """See the builtin zip_ Difference from stdlib: accepts ``longest`` kwarg, which makes this method act like itertools.zip_longest_ (also accepts ``fill_value``). ``strict=True`` is only supported in python >=3.10 and will raise a TypeError in lower versions. .. _zip: https://docs.python.org/3/library/functions.html#zip .. _zip_longest: https://docs.python.org/3/library/itertools.html#itertools.zip_longest """ # noqa iterables = [self] + list(iterables) if longest: return FIt.zip_longest(*iterables, fill_value=fill_value) try: length = min(len_or_none(it) for it in iterables) except TypeError: length = None if strict: inner = zip(*iterables, strict=True) else: inner = zip(*iterables) return FIt(inner, length)
def cycle(self, n: Optional[int] = None): """See itertools.cycle_ Difference from stdlib: accepts ``n`` argument, for how many times it should be repeated. .. _itertools.cycle: https://docs.python.org/3/library/itertools.html#itertools.cycle """ # noqa this_len = len_or_none(self) if n is None: n = float("inf") elif this_len is not None: length = n * this_len return FIt(cycle(self), length).islice(length) def gen(it): cached = [] for item in it: cached.append(item) yield item count = 1 while count < n: yield from cached count += 1 return FIt(gen(self))
def interleave(self, *iterables): """Interleave items from any number of iterables When an iterable is exhausted, items continue to be yielded from the remaining iterables. :param iterables: iterables providing items :return: """ iterables = [self] + list(iterables) try: length = sum(len_or_none(it) for it in iterables) except TypeError: length = None def gen(its): its = [iter(it) for it in its] while its: next_its = [] for it in its: try: yield next(it) next_its.append(it) except StopIteration: pass its = next_its return FIt(gen(iterables), length)
def islice(self, start: int, stop: Optional[int] = None, step: Optional[int] = None) -> FIt[T]: """See itertools.islice_ Difference from stdlib: if the FIt length is known, negative indices are allowed, although a negative step size is still not. Consumes elements up to and including the start index; subsequent items are consumed by the returned iterator. .. _itertools.islice: https://docs.python.org/3/library/itertools.html#itertools.islice """ # noqa if stop is None: stop = start start = 0 if step is None: step = 1 this_len = len_or_none(self) if this_len is None: if start < 0 or stop < 0: raise ValueError(neg_idx_msg) length = None else: if start < 0: start = max(this_len + start, 0) if stop < 0: stop = max(this_len + stop, 0) length = max(math.ceil((min(stop, this_len) - start) / step), 0) return FIt(islice(self, start, stop, step), length)
def tail(self, n: int): """Return an iterator over the last ``n`` items""" this_len = len_or_none(self) if this_len is None: length = None else: length = min(n, this_len) return FIt(deque(self, maxlen=n), length)
def assert_equal_it(it1, it2): len1 = len_or_none(it1) len2 = len_or_none(it2) if len1 is not None and len2 is not None: assert len1 == len2 count = 0 for item1, item2 in zip_longest(it1, it2): assert item1 == item2 count += 1 if len1 is not None: assert count == len1 if len2 is not None: assert count == len2
def permutations(self, r: Optional[int] = None): """See itertools.permutations_ .. _itertools.permutations: https://docs.python.org/3/library/itertools.html#itertools.permutations """ # noqa this_len = len_or_none(self) if this_len is None: length = None else: length = n_permutations(this_len, r) return FIt(permutations(self, r), length)
def combinations_with_replacement(self, r: int): """See itertools.combinations_with_replacement_ .. _itertools.combinations_with_replacement: https://docs.python.org/3/library/itertools.html#itertools.combinations_with_replacement """ # noqa this_len = len_or_none(self) if this_len is None: length = None else: length = nCr(this_len, r, True) return FIt(combinations_with_replacement(self, r), length)
def chain(self, *iterables): """See itertools.chain_ .. _itertools.chain: https://docs.python.org/3/library/itertools.html#itertools.chain """ # noqa iterables = [self] + list(iterables) lengths = [len_or_none(it) for it in iterables] try: length = sum(lengths) except TypeError: length = None return FIt(chain.from_iterable(iterables), length)
def zip_longest(self, *iterables, fill_value=None): """See itertools.zip_longest_ .. _itertools.zip_longest: https://docs.python.org/3/library/itertools.html#itertools.zip_longest """ # noqa iterables = [self] + list(iterables) try: length = max(len_or_none(it) for it in iterables) except TypeError: length = None return FIt(zip_longest(*iterables, fillvalue=fill_value), length)
def islice(self, start: int, stop: Optional[int] = None, step: int = 1): """See itertools.islice_ .. _itertools.islice: https://docs.python.org/3/library/itertools.html#itertools.islice """ # noqa if stop is None: stop = start start = 0 this_len = len_or_none(self) if this_len is None: length = None else: length = max(math.ceil((min(stop, this_len) - start) / step), 0) return FIt(islice(self, start, stop, step), length)
def progress(self, **kwargs): """Create a tqdm progress bar for this iterable. :param kwargs: passed to tqdm instance :return: FIt wrapping a tqdm instance """ try: from tqdm import tqdm return FIt(tqdm(self, **kwargs)) except ImportError: warnings.warn("Progress bar is not available: pip install tqdm") this_len = len_or_none(self) if this_len is None: this_len = kwargs.get("total") return FIt(self, this_len)
def __getitem__(self, idx: Union[int, slice]): """Shorthand for ``.get`` or ``.islice``. Depending on whether an integer or slice is given, returns either a single item (see ``.get``), or an ``FIt`` instance (see ``.islice``). Consumes necessary iterator elements. """ if isinstance(idx, int): if idx < 0 and len_or_none(self) is None: raise ValueError(neg_idx_msg) return self.get(idx) elif isinstance(idx, slice): return self.islice(idx.start, idx.stop, idx.step) else: raise TypeError("FIt indices must be integers or slices, not " + type(idx).__name__)
def combinations(self, r: int, replace=False): """See itertools.combinations_ Difference from stdlib: ``replace`` argument to use itertools.combinations_with_replacement_ .. _itertools.combinations: https://docs.python.org/3/library/itertools.html#itertools.combinations .. _itertools.combinations_with_replacement: https://docs.python.org/3/library/itertools.html#itertools.combinations_with_replacement """ # noqa if replace: return FIt.combinations_with_replacement(self, r) this_len = len_or_none(self) if this_len is None: length = None else: length = nCr(this_len, r) return FIt(combinations(self, r), length)
def sliding_window(self, n: int): """Iterate over ``n``-length tuples forming a sliding window over the iterable. :param n: window size :return: FIt of tuples """ this_len = len_or_none(self) length = None if this_len is None else max(this_len - n + 1, 0) def gen(it): window = deque(it.take(n), maxlen=n) if len(window) < n: return yield tuple(window) for item in it: window.append(item) yield tuple(window) return FIt(gen(self), length)
def zip(self, *iterables, longest=False, fill_value=None): """See the builtin zip_ Difference from stdlib: accepts ``longest`` kwarg, which makes this method act like itertools.zip_longest_ (also accepts ``fill_value``) .. _zip: https://docs.python.org/3/library/functions.html#zip .. _zip_longest: https://docs.python.org/3/library/itertools.html#itertools.zip_longest """ # noqa iterables = [self] + list(iterables) if longest: return FIt.zip_longest(*iterables, fill_value=fill_value) try: length = min(len_or_none(it) for it in iterables) except TypeError: length = None return FIt(zip(*iterables), length)
def __init__(self, iterable: Iterable, length=None): """Iterator class providing many postfix functional methods. Most of these methods can also be used as static methods which take any iterable as the first argument. Where possible, the returned FIt instances have a length, which expresses how many items *remain* in the iterator. Where possible, all iteration is evaluated lazily. :param iterable: iterable to wrap :param length: explicitly provide a length if you know it but the iterable doesn't """ # noqa self.iterator = iter(iterable) if length is None: self.init_length = len_or_none(iterable) else: self.init_length = int(length) self.consumed = 0 self.cache = deque()
def flatten(self, levels: Optional[int] = None, split_strings: bool = False): """Recursively flatten arbitrary iterables (depth-first). By default, strings are not treated as iterables to descend into. If ``split_strings`` is truthy, their characters will be yielded individually. :param levels: How deep in the iterable to flatten (default all levels) :param split_strings: Whether to yield individual characters from strings (default False) :return: FIt """ # noqa if levels is None: levels = float("inf") def gen(obj, lvls): if isinstance(obj, str) and (len(obj) == 1 or not split_strings): yield obj return try: it = iter(obj) except TypeError: yield obj return if lvls <= 0: yield from it else: for item in it: yield from gen(item, lvls - 1) this_len = len_or_none(self) if levels == 0 and this_len is not None: length = this_len else: length = None return FIt(gen(self, levels), length)
def chunk(self, chunksize: int): """Iterate over ``chunksize``-or-shorter lists which are chunks of the iterable. :param chunksize: maximum length for each chunk (all but the last chunk will be this size) :return: FIt of lists """ # noqa this_len = len_or_none(self) length = None if this_len is None else math.ceil(this_len / chunksize) def gen(it): it = iter(it) taken = [] while True: for _ in range(chunksize): try: taken.append(next(it)) except StopIteration: yield taken return yield taken taken = [] return FIt(gen(self), length)
def get(self, n: int, default=EMPTY) -> T: """Alias for ``FIt.nth``: returns the nth item or a default value. If default is not given, raises IndexError. Accepts negative index if length is known. Consumes elements up to and including the given index. Cannot be safely used as a static method. """ if n < 0: this_len = len_or_none(self) if this_len is None: raise ValueError(neg_idx_msg) n = this_len + n if n < 0: raise IndexError("FIt index out of range") try: result = next(self.islice(n, n + 1)) except StopIteration: if default is EMPTY: raise IndexError("FIt index out of range") else: result = default return result
def starmap(self, function: Callable): """See itertools.starmap_ .. _itertools.starmap: https://docs.python.org/3/library/itertools.html#itertools.starmap """ # noqa return FIt(starmap(function, self), len_or_none(self))
def test_cycle(fiter, lst): cycling = fiter.cycle(3) if len_or_none(fiter): assert len(cycling) == len(lst) * 3 assert_equal_it(cycling, lst * 3)
def enumerate(self, start=0): """See the builtin enumerate_ .. _enumerate: https://docs.python.org/3/library/functions.html#enumerate """ return FIt(enumerate(self, start), len_or_none(self))
def test_len_or_none(lst): length = len(lst) assert len_or_none(lst) == length assert len_or_none(iter(lst)) is None