real-time decorator for functions and generators

问题

I have a situation in which I need to hook certain functions so that I can inspect the return values and track them. This is useful for tracking for example running averages of values returned by methods/functions. However, these methods/function can also be generators.

However, if i'm not wrong, python detects generators when parsing and when the function is called at runtime it always returns a generator. Thus I can't simply do something like:

import types
def decorator(func):
    average = None # assume average can be accessed by other means
    def wrap(*args, **kwargs):
        nonlocal average
        ret_value = func(*args, **kwargs)
        #if False wrap is still a generator 
        if isinstance(ret_value, types.GeneratorType): 
           for value in ret_value:
              # update average
              yield value
        else:
            # update average
            return ret_value # ret_value can't ever be fetched
    return wrap

And yielding in this decorator is necessary, since I need to track the values as the caller iterates this decorated generator (i.e. "real-time"). Meaning, I can't simply replace the for and yield with values = list(ret_value), and return values. (i.e.) If the func is a generator it needs to remain a generator once decorated. But if func is a pure function/method, even if the else is executed, wrap still remains a generator. Meaning, the ret_value can't ever be fetched.

A toy example of using such a generator would be:

@decorated
def some_gen(some_list):
    for _ in range(10):
       if some_list[0] % 2 == 0:
           yield 1
       else:
           yield 0
def caller():
   some_list = [0]
   for i in some_gen(some_list):
      print(i)
      some_list[0] += 1 # changes what some_gen yields

For the toy example, there may be simpler solutions, but it's just to prove a point.

Maybe I'm missing something obvious, but I did some research and didn't find anything. The closest thing I found was this. However, that still doesn't let the decorator inspect every value returned by the wrapped generator (just the first). Does this have a solution, or are two types of decorators (one for functions and one for decorators) necessary?

回答1:

Once solution I realized is:

def as_generator(gen, avg_update):
     for i in gen:
         avg_update(i)
         yield i

import types
def decorator(func):
    average = None # assume average can be accessed by other means
    def wrap(*args, **kwargs):
        def avg_update(ret_value):
            nonlocal average
            #update average
            pass

        ret_value = func(*args, **kwargs)
        #if False wrap is still a generator 
        if isinstance(ret_value, types.GeneratorType): 
           return as_generator(ret_value, avg_update)
        else:
            avg_update(ret_value)
            return ret_value # ret_value can't ever be fetched
    return wrap

I don't know if this is the only one, or if there exists one without making a separate function for the generator case.

来源：https://stackoverflow.com/questions/53686856/real-time-decorator-for-functions-and-generators

标签

python-3.x

generator

python-decorators