Skip to content

Instantly share code, notes, and snippets.

@krackers
Forked from hzhu212/python_stream.md
Created August 25, 2023 09:22
Show Gist options
  • Save krackers/011fc68ed1d665099a397f9d8e658f47 to your computer and use it in GitHub Desktop.
Save krackers/011fc68ed1d665099a397f9d8e658f47 to your computer and use it in GitHub Desktop.
Python iterator with chain-call like Java stream

stream.py:

import itertools

class Stream(object):
    """
    a stream object for chain-call
    """
    def __init__(self, iterable):
        self._iterable = iterable

    def __iter__(self):
        return self

    def __next__(self):
        return next(self._iterable)

    def skip(self, n):
        self._iterable = itertools.islice(self._iterable, n, None, 1)
        return self

    def limit(self, n):
        self._iterable = itertools.islice(self._iterable, None, n, 1)
        return self

    def step(self, n):
        self._iterable = itertools.islice(self._iterable, None, None, n)
        return self

    def __getitem__(self, slc):
        if not isinstance(slc, slice):
            raise ValueError(f'"[]" operator supports only slice object, but {slc!r} got')
        self._iterable = itertools.islice(self._iterable, slc.start, slc.stop, slc.step)
        return self

    def apply(self, fun, *args, **kwargs):
        self._iterable = (fun(x, *args, **kwargs) for x in self._iterable)
        return self

    def tolist(self):
        return list(self)

Use cases:

In [15]: Stream(range(100)).skip(10).step(3).limit(10).apply(lambda x: x//3).tolist()                                                                                       
Out[15]: [3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

In [2]: Stream(range(100))[10:20:3].apply(lambda x: x+1).tolist()                                                                                                          
Out[2]: [11, 14, 17, 20]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment