Welcome, guest | Sign In | My Account | Store | Cart
NOTE: Recipes have moved! Please visit GitHub.com/activestate/code for the current versions.

This recipe enables the use of the yield statement within a method by decorating that method with a wrapper for a generator object. The purpose of using this decorator is to allow the method to be invoked using the normal calling syntax. A caller need not know the method is actually a generator and can focus solely on the method's interface rather than how it is implemented.

Python, 170 lines
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
#gendec.py

import weakref

class MethodGeneratorProxy(object):
  '''
  Wraps a generator method in a callable. On each call to an instance of this
  class, one of the following occurs: the generator advances one iteration and
  its result is returned, the generator is reset and None is returned, or the
  generator throws a StopIteration (which is caught) and None is returned.
  The generator is automatically re-instantiated on the next call after the
  StopIteraction exception is raised.

  This class does not maintain a strong reference to the object to which the
  method is attached. It will not impede garbage collection.

  @ivar func: Method that, when called, instantiates a generator
  @type func: function
  @ivar gen: Instantiated generator
  @type gen: generator
  @ivar reset_flag: Param to look for in keyword args as the signal to reset
  @type reset_flag: string
  '''     
  def __init__(self, func, reset_flag):
    '''
    Initializes an instance.

    See instance variables for parameter descriptions.
    '''       
    self.func = func
    self.gen = None
    self.reset_flag = reset_flag

  def __call__(self, obj, *args, **kwargs):
    '''
    Generate the next item or reset the generator.

    @param obj: Object to which the generator is attached
    @type obj: object
    @param args: Positional arguments to the generator method
    @type args: list
    @param kwargs: Keyword arguments to the generator method
    @type kwargs: dictionary
    '''       
    if kwargs.get(self.reset_flag):
      # reset the generator and return None
      self.gen = None
      return None
    elif self.gen is None:
      # create a new generator
      try:
        wobj = weakref.proxy(obj)
      except TypeError:
        wobj = obj
      self.gen = self.func(wobj, *args, **kwargs)

    try:
      # generate next item
      return self.gen.next()
    except StopIteration:
      # destroy the generator and return None
      self.gen = None
      return None

def generator_method(cls_name, reset_flag='reset_gen'):
  '''
  Decorator for methods that act as generators. Methods decorated as such can
  be called like normal methods, but can (and should) include yield statements.
  The parameters passed to the first invocation of the method are used for all
  subsequent calls until the generator is reset.

  Generator methods can be overriden in subclasses as long as the name provided
  is unique to each class in the inheritence tree (i.e. make it the name of the
  class and everything will work fine.

  To reset a generator method, pass True in a keyword argument with the name
  specified in reset_flag. Generator methods in parent classes must be reset
  explicitly (i.e. Parent.MethodName(self, reset_gen=True).

  @param cls_name: Name unique to the inheritence tree of this class
  @type cls_name: string
  @param reset_flag: Name of a parameter that will reset the generator
  @type reset_flag: string
  '''     
  # define another function that takes just the function as an argument
  # we must do this to deal with the fact that we need the name argument above
  def generator_method_internal(func):
    # build a name for the method generator
    name = '_%s_%s_gen_proxy_' % (cls_name, func.func_name)
    # define a replacement for a method that calls the generator instead
    def generator_method_invoke(obj, *args, **kwargs):
      try:
        # try to get a generator defined for the called method
        gen = getattr(obj, name)
      except AttributeError:
        # build a new generator for the called method
        gen = MethodGeneratorProxy(func, reset_flag)
        setattr(obj, name, gen)
      # call the generator and return its result
      return gen(obj, *args, **kwargs)
    # return our wrapping for the method
    return generator_method_invoke
  # return the true decorator for the method
  return generator_method_internal



# test.py
from gendec import generator_method

class Test(object):
  def __del__(self):
    print '* Test instance freed'

  def Reset(self):
    self.GetWords(reset_gen=True)

  @generator_method('Test')
  def GetWords(self):
    yield 'the quick'
    yield 'brown fox'
    yield 'jumped over'
    yield 'the lazy'
    yield 'dog'

class Foo(Test):
  def __del__(self):
    print '* Foo instance freed'

  def ResetAll(self):
    super(Foo, self).GetWords(reset_gen=True)
    self.GetWords(reset_gen=True)

  @generator_method('Foo')
  def GetWords(self):
    i = 0
    while 1:
      i += 1
      s = super(Foo, self).GetWords()
      yield '<%d> %s' % (i,s)

print '*** Test instance ***'
t = Test()
for i in range(8):
  print t.GetWords()
print '*** Resetting'
t.Reset()
for i in range(3):
  print t.GetWords()

print
print '*** Test instance #2 ***'
s = Test()
for i in range(3):
  print s.GetWords()

print
print '*** Foo subclass instance ***'
f = Foo()
for i in range(8):
  print f.GetWords()
print '*** Resetting Foo method only'
f.Reset()
for i in range(5):
  print f.GetWords()
print '*** Resetting Foo and Test methods'
f.ResetAll()
for i in range(3):
  print f.GetWords()
print

I came across a situation recently where I wanted a method in one of my classes to return a different piece of information depending on the state of the class. Had the number of states been small, using an instance variable to track the state and then an if/else block in the method would have worked easily enough. However, the number of states was quite large and the transition triggers depended upon a number of factors.

A generator method, a method of a class instance that allows the use of the yield statement but is invoked like a normal method, is a simple solution in this circumstance. The method saves state by design and alleviates the need to track state in explicit instance variables. The recipe below uses a decorator and a generator proxy to define generator methods which can be subclassed, can be reset at any time, and do not inhibit garbage collection.

4 comments

Paul Moore 12 years, 10 months ago  # | flag

I'm not entirely sure I see the point. You can get 90% of the effects by directly capturing the next() method of the generator, like so:

>>> class T:
...     def foo(self):
...         yield 1
...         yield 2
...         yield 3
...
>>> t = T()
>>> f = t.foo().next
>>> f()
1
>>> f()
2
>>> f()
3
>>> f()
Traceback (most recent call last):
  File "", line 1, in ?
StopIteration
>>> f = t.foo().next
>>> f()
1

You get a StopIteration exception rather than None and an automatic reset, and the reset is a bit more explicit, but it seems a lot clearer to me.

Of course, clarity is very personal, so if you like this form, that's fine.

Peter Parente (author) 12 years, 10 months ago  # | flag

Encapsulation. The real purpose of this recipe is to hide the fact that a method is actually a generator. A client object invoking a method probably should not have to worry about whether the method can be called directly or if next should be invoked instead. This recipe allows the caller to use the normal method invocation syntax without knowing anything about that method's implementation.

Daniel Serodio 12 years, 8 months ago  # | flag

What about generator functions. Nice recipe. How can I wrap a function instead of a method?

Q Neill 8 years, 2 months ago  # | flag

@Daniel: Why wrap a function when you can just yield directly from it?

Created by Peter Parente on Mon, 21 Mar 2005 (PSF)
Python recipes (4591)
Peter Parente's recipes (1)

Required Modules

Other Information and Tasks