7
\$\begingroup\$

I'm not a Python developper, but I enjoy programming with it, and for a project I wanted to have generators that I can easily index. Using python's slice model is obviously the way to go, and here's the solution I've come up with.

class _SubscriptableGenerator():
    def __init__(self, generator, *args):
        self.gen = generator(*args)

    def __getitem__(self, key):
        try:
            if isinstance(key, int):
                self.ignore(key)
                yield next(self.gen)
            else:
                step = key.step if key.step else 1
                start = key.start if key.start else 0

                i = start
                self.ignore(start)

                while i < key.stop:
                    yield next(self.gen)
                    i = i + step
                    self.ignore(step-1)
        except Exception:
            self.raiseInvalidSlice(key)

    def raiseInvalidSlice(self, key):
        raise KeyError("{0} is not a valid key (only int and [x:y:z] slices are implemented.".format(key))

    def ignore(self, n):
        for i in range(n):
            next(self.gen)

It is not intended to be called by the user of the module, that's internal code. I for example define my generators like so

def _myGen(arg1, arg2):
    while 1:
        yield something

and provide them wrapped in my class

def myGen(*args):
    return _SubscriptableGenerator(_myGen, *args)

I'd like to know what a more pythonic solution would be, if there are things to fix, etc. I am not sure about the way to handle exceptions on the key.

\$\endgroup\$
9
\$\begingroup\$

The most Pythonic solution would be to use itertools.islice from the standard library. For example, like this:

from itertools import islice

class Sliceable(object):
    """Sliceable(iterable) is an object that wraps 'iterable' and
    generates items from 'iterable' when subscripted. For example:

        >>> from itertools import count, cycle
        >>> s = Sliceable(count())
        >>> list(s[3:10:2])
        [3, 5, 7, 9]
        >>> list(s[3:6])
        [13, 14, 15]
        >>> next(Sliceable(cycle(range(7)))[11])
        4
        >>> s['string']
        Traceback (most recent call last):
            ...
        KeyError: 'Key must be non-negative integer or slice, not string'

    """
    def __init__(self, iterable):
        self.iterable = iterable

    def __getitem__(self, key):
        if isinstance(key, int) and key >= 0:
            return islice(self.iterable, key, key + 1)
        elif isinstance(key, slice):
            return islice(self.iterable, key.start, key.stop, key.step)
        else:
            raise KeyError("Key must be non-negative integer or slice, not {}"
                           .format(key))

Note that I've given the class a better name, written a docstring, and provided some doctests.

\$\endgroup\$
  • \$\begingroup\$ Nice, I didn't know about this itertools feature. Thanks ! \$\endgroup\$ – teh internets is made of catz Oct 22 '13 at 15:05

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.