Skip to content

Instantly share code, notes, and snippets.

@mjwillson
Last active Aug 29, 2015
Embed
What would you like to do?
Decorate a generator function (or other iterator-returning function) as a multi-shot iterable. A fix for many Python gotchas relating to use of one-shot iterators
class iterable(object):
"""Decorates a generator function (or any other iterator-returning
function) as something which implements the iterable protocol and
can be safely passed to other code which may iterate over it
multiple times.
Usage:
@iterable
def foo():
yield 1
yield 2
works like a plain generator function if you really want:
generator = foo()
[x for x in generator]
=> [1, 2]
But iterating a generator a second time has the following, silent
and often undesired result, which has been the source of many
subtle, irritating bugs for me:
[x for x in generator]
=> []
Instead we can now iterate the wrapped generator function
directly, as many times as we like:
[x for x in foo]
=> [1, 2]
[x for x in foo]
=> [1, 2]
This is useful if you want to pass a generic multiply-iterable
object to someone, but have it be lazy rather than having to
convert to a list first.
It works for methods too:
class Foo():
@iterable
def foo(self):
yield self
yield self
it = Foo().foo
[x for x in it]
=> [<__main__.Foo instance at 0x7f2144e2ce18>,
<__main__.Foo instance at 0x7f2144e2ce18>]
[x for x in it]
=> [<__main__.Foo instance at 0x7f2144e2ce18>,
<__main__.Foo instance at 0x7f2144e2ce18>]
It'd be nice if we could do something similar for the more concise
generator expression syntax, although sadly this syntax inherently
produces one-shot iterators. You can write a function that returns
a generator expression (or any other iterator) though, e.g.:
@iterable
def my_iterable():
return (expensive_transform(x) for x in [1,2,3])
You can also specify arguments to pass to the iterator-returning
function, which makes more sense if you're using it as a plain
function rather than a decorator, and is particularly useful with
higher-order functions like itertools.imap (python 2) or map
(python 3), which annoyingly return one-shot iterators even if
their argument is multiply iterable:
result = map(lambda x : x+1, [1,2,3])
[x for x in result]
=> [2,3,4]
[x for x in result]
=> [] # Gah!!!
The cure:
result = iterable(map, lambda x : x+1, [1,2,3])
[x for x in result]
=> [2,3,4]
[x for x in result]
=> [2,3,4]
"""
def __init__(self, func, *args, **kwargs):
self.func = func
self.args = args
self.kwargs = kwargs
def __call__(self):
"""Proxy through standard function calls to the function it wraps"""
return self.func(*self.args, **self.kwargs)
"""But also expose the iterable protocol"""
__iter__ = __call__
def __get__(self, obj, objtype=None):
"""Make method binding work like for a normal function"""
return iterable(self.func.__get__(obj, objtype))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment