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.
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.
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:
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.
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.
What about generator functions. Nice recipe. How can I wrap a function instead of a method?
@Daniel: Why wrap a function when you can just yield directly from it?