钩住生成器函数的最佳方法(如果有的话)

时间:2016-03-07 10:05:50

标签: python python-3.x generator python-decorators

我正在写一个非常简单的装饰器来给我一些关于函数的基本调试信息。

from functools import wraps
from time import perf_counter

class debug(object):
    def __init__(self, Time=False, Parameters=False, Doc=False):
        self.t = Time
        self.p = Parameters
        self.d = Doc

    def __call__(self, func):
        @wraps(func)
        def run(*args, **kwargs):
            params = ""
            if self.p:
                params = ", ".join(["{}".format(arg) for arg in args] + ["{}={}".format(k, v) for k, v in kwargs.items()])
            print("\n\tDebug output for '{}({})'".format(func.__name__, params))
            if self.d:
                print('\tDocstring: "{}"'.format(func.__doc__))
            if self.t:
                t1 = perf_counter()
            val = func(*args, **kwargs)
            if self.t:
                t2 = perf_counter()
                print("\tTime Taken: {:.3e} seconds".format(t2 - t1))
            print("\tReturn Type: '{}'\n".format(type(val).__name__))
            return val
        return run

这对于正常功能来说都很好。

@debug(Parameters=True, Time=True, Doc=True)
def foo(i, j=5):
    """Raises i to 2j"""
    for _ in range(j):
        i **= 2
    return i

i = foo(5, j=3)
# Output:
"""
    Debug output for 'foo(5, j=3)'
    Docstring: "Raises i to 2j"
    Time Taken: 1.067e-05 seconds
    Return Type: 'int'
"""

然而,发电机是一个不同的故事。

@debug(Parameters=True, Time=True, Doc=True)
def bar(i, j=2):
    """Infinite iterator of increment j"""
    while True:
        yield i
        i += j

b = bar()  # Output occurs here
next(b) # No output

现在,根据我编写的内容,这是完全可以预料到的,但我想知道如何挂钩.__next__()方法或者最好的解决方法是什么。

2 个答案:

答案 0 :(得分:1)

如果将生成器作为输入(在文件顶部添加__call__),您只需更改import types方法并返回生成器:

def __call__(self, f):
    if isinstance(f, types.GeneratorType):
        def run_gen(*args, **kwargs):
            # do pre stuff...
            for _ in f(*argw, **kwargs):
                yield _
            # do post stuff...
        return run_gen
    else:
        def run(*args, **kwargs):
            # do pre stuff...
            r = f(*argw, **kwargs)
            # do post stuff...
            return r
        return run

答案 1 :(得分:1)

您无法替换function.next,因为它是只读值。但你可以这样做(参见debug_generator函数):

来自functools导入包装的

    导入检查

class debug(object):
    def __init__(self, Time=False, Parameters=False, Doc=False):
        self.t = Time
        self.p = Parameters
        self.d = Doc

    def __call__(self, func):
        @wraps(func)
        def debug_generator(func):
            for i, x in enumerate(list(func)):
                # here you add your debug statements
                print "What you want: step %s" % i
                yield x
        @wraps(func)
        def run(*args, **kwargs):
            params = ""
            if self.p:
                params = ", ".join(["{}".format(arg) for arg in args] + ["{}={}".format(k, v) for k, v in kwargs.items()])
            print("\n\tDebug output for '{}({})'".format(func.__name__, params))
            if self.d:
                print('\tDocstring: "{}"'.format(func.__doc__))
            val = func(*args, **kwargs)
            print("\tReturn Type: '{}'\n".format(type(val).__name__))

            if inspect.isgenerator(val):
                return debug_generator(val)
            return val
        return run

基本上你只是从想要调试的生成器中获取所有值,然后再次yield它们,在循环中添加调试语句。