HiveBrain v1.2.0
Get Started
← Back to all entries
patternpythonMinor

Subscriptable/Indexable generator

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
indexablegeneratorsubscriptable

Problem

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.

Solution

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.

Code Snippets

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))

Context

StackExchange Code Review Q#33060, answer score: 9

Revisions (0)

No revisions yet.