Last active
August 29, 2015 14:27
-
-
Save mjwillson/805685de34bcad1e25cc to your computer and use it in GitHub Desktop.
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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