Wrapping generator functions in Python

Question

I am writing some code that traverses a structure that may have cyclic references. Rather than explicitly doing checks at the beginning of the recursive functions I thought that I would create a decorator that didn't allow a function to be called more than once with the same arguments.

Below is what I came up with. As it is written, this will try to iterate over Nonetype and raise an exception. I know that I could fix it by returning say an empty list, but I wanted to be more elegant. Is there a way to tell from within the decorator whether the function being decorated is a generator function or not? This way I could conditionally raise StopIteration if it is a generator or just return None otherwise.

previous = set()
def NO_DUPLICATE_CALLS(func):
    def wrapped(*args, **kwargs):
        if args in previous:
            print 'skipping previous call to %s with args %s %s' % (func.func_name, repr(args), repr(kwargs))
            return
        else:
            ret = func(*args, **kwargs)
            previous.add(args)
            return ret
    return wrapped

@NO_DUPLICATE_CALLS
def foo(x):
    for y in x:
        yield y

for f in foo('Hello'):
    print f

for f in foo('Hello'):
    print f

The `inspect` module is actually pretty neat, when you think about it. — JAB, Jun 20 '11 at 21:18

JAB · Accepted Answer · 2011-06-20T19:49:22.427

5

Okay, check this out:

>>> from inspect import isgeneratorfunction
>>> def foo(x):
...    for y in x:
...        yield y
...
>>> isgeneratorfunction(foo)
True

This requires Python 2.6 or higher, though.

edited Jun 20 '11 at 19:49

answered Jun 20 '11 at 16:56

JAB

20,783
6
71
80

1

>>> isinstance(s, generator) Traceback (most recent call last): File "", line 1, in NameError: name 'generator' is not defined – Mariy Jun 20 '11 at 17:01
@JAB: Is "generator" a python 3.2 thing? "isinstance(func, generator)" just gives a NameError in Python 2.7 – Gerrat Jun 20 '11 at 17:05
1

@cldy, Gerrat: Sorry, just noticed that even though the class name is 'generator', you have to import `GeneratorType` from the `types` module and use that for it to work (updated my answer to reflect that). It's silly, though, because things like `isinstance(1, int)` work just fine. (Well, in Python 2.2+.) I guess it's due to how `int` and such are built-in functions as well, but `generator` is not. – JAB Jun 20 '11 at 17:12
1

The function should not raise a StopIteration, it should return an iterator that raises a StopIteration on the first `next` call, an empty list would work fine. – Andrew Clark Jun 20 '11 at 17:25
Also, this just won't work. Generator functions are functions that return a generator when they are called, the function itself will not be an instance of the `generator` class. You would have to use `isinstance(func(), GeneratorType)` which defeats the purpose. – Andrew Clark Jun 20 '11 at 17:30
Redid my answer to show a nice way of checking if it's a generator function. – JAB Jun 20 '11 at 19:49

score 4 · Answer 2 · edited May 23 '17 at 12:20

Unfortunately there is not really a good way to know whether a function might return some type of iterable without calling it, see this answer to another question for a pretty good explanation of some potential issues.

However, you can get around this by using a modified memoize decorator. Normally memoizing decorators would create a cache with the return values for previous parameters, but instead of storing the full value you could just store the type of the return value. When you come across parameters you have already seen just return a new initialization of that type, which would result in an empty string, list, etc.

Here is a link to memoize decorator to get you started:
http://wiki.python.org/moin/PythonDecoratorLibrary#Memoize

Deleted my answer as it was ultimately useless. Looks like searching the results of `dis.dis(foo)` for `YIELD_VALUE`, as mentioned in that question you linked, might be the best bet. On a side note, eric was technically using a form of memoization already with his `previous` set. — JAB, Jun 20 '11 at 18:00

Wrapping generator functions in Python

2 Answers2