检查函数是否被称为装饰器

Question

在以下最小示例中，decorate 被调用了两次。首先使用@decorate，其次是普通函数调用decorate(bar)。

def decorate(func):
    print(func.__name__)
    return func

@decorate
def bar():
    pass

decorate(bar)

如果调用是通过使用@decorate 或作为普通函数调用调用的，是否可以看到decorate 的内部？

Answer 1

@decorator 语法只是语法糖，因此两个示例具有相同的行为。这也意味着您在它们之间所做的任何区分都可能没有您想象的那么有意义。

尽管如此，您可以使用inspect 来阅读您的脚本并查看在上面的框架中是如何调用装饰器的。

import inspect

def decorate(func):
    # See explanation below
    lines = inspect.stack(context=2)[1].code_context
    decorated = any(line.startswith('@') for line in lines)

    print(func.__name__, 'was decorated with "@decorate":', decorated)
    return func

请注意，我们必须将 context=2 指定给 inspect.stack 函数。 context 参数表示必须返回当前行周围的代码行数。在某些特定情况下，例如在装饰子类时，当前行位于类声明而不是装饰器上。 The exact reason for this behaviour has been explored here.

示例

@decorate
def bar():
    pass

def foo():
    pass
foo = decorate(foo)

@decorate
class MyDict(dict):
    pass

输出

bar was decorated with "@decorate": True
foo was decorated with "@decorate": False
MyDict was decorated with "@decorate": True

警告

还有一些我们难以克服的极端情况，例如装饰器和类声明之间的换行符。

# This will fail
@decorate

class MyDict(dict):
    pass

Answer 2

Olivier 的回答打消了我的想法。但是，由于inspect.stack() 是一个特别昂贵的电话，我会考虑选择使用以下方式：

frame = inspect.getframeinfo(inspect.currentframe().f_back, context=1)
if frame.code_context[0][0].startswith('@'): 
    print('Used as @decorate: True')
else:
    print("Used as @decorate: False")

Answer 3

与普遍认为的相反，@decorator 和 decorator(…) 并不完全相同。第一个运行 before 名称绑定，后者运行 after 名称绑定。对于顶级函数的常见用例，这允许以低成本测试适用哪种情况。

import sys

def decoraware(subject):
    """
    Decorator that is aware whether it was applied using `@deco` syntax
    """
    try:
        module_name, qualname = subject.__module__, subject.__qualname__
    except AttributeError:
        raise TypeError(f"subject must define '__module__' and '__qualname__' to find it")
    if '.' in qualname:
        raise ValueError(f"subject must be a top-level function/class")
    # see whether ``subject`` has been bound to its module
    module = sys.modules[module_name]
    if getattr(module, qualname, None) is not subject:
        print('@decorating', qualname)  # @decoraware
    else:
        print('wrapping()', qualname)   # decoraware()
    return subject

这个例子只会打印它是如何应用的。

>>> @decoraware
... def foo(): ...
...
@decorating foo
>>> decoraware(foo)
wrapping() foo

不过，可以使用相同的方法在每个路径中运行任意代码。

如果应用了多个装饰器，您必须决定是要顶部还是底部主题。对于顶级功能，代码无需修改即可工作。对于底部主题，在检测前使用subject = inspect.unwrap(subject) 解包。

---

可以在 CPython 上以更通用的方式使用相同的方法。使用 sys._getframe(n).f_locals 可以访问应用了装饰器的本地命名空间。

def decoraware(subject):
    """Decorator that is aware whether it was applied using `@deco` syntax"""
    modname, topname = subject.__module__, subject.__name__
    if getattr(sys.modules[modname], topname, None) is subject:
        print('wrapping()', topname, '[top-level]')
    else:
        at_frame = sys._getframe(1)
        if at_frame.f_locals.get(topname) is subject:
            print('wrapping()', topname, '[locals]')
        elif at_frame.f_globals.get(topname) is subject:
            print('wrapping()', topname, '[globals]')
        else:
            print('@decorating', topname)
    return subject

---

请注意，与 pickle 类似，如果主题的 __qualname__/__name__ 被篡改或从其定义的命名空间中被 del'ed，则此方法将失败。

Answer 4

在前两个答案的基础上，我编写了一个通用函数，它应该在几乎所有实际情况下都能按预期工作。我使用 Python 3.6、3.7 和 3.8 对其进行了测试。

在将此函数复制粘贴到您的代码中之前，请确保您最好改用decorator module。

def am_I_called_as_a_decorator(default=False):
    """This function tries to determine how its caller was called.

    The value returned by this function should not be blindly trusted, it can
    sometimes be inaccurate.

    Arguments:
        default (bool): the fallback value to return when we're unable to determine
                        how the function was called

    >>> def f(*args):
    ...     if am_I_called_as_a_decorator():
    ...         print("called as decorator with args {!r}".format(args))
    ...         if len(args) == 1:
    ...             return args[0]
    ...         return f
    ...     else:
    ...         print("called normally with args {!r}".format(args))
    ...
    >>> f()
    called normally with args ()
    >>> @f                              #doctest: +ELLIPSIS
    ... def g(): pass
    ...
    called as decorator with args (<function g at ...>,)
    >>> @f()
    ... class Foobar: pass
    ...
    called as decorator with args ()
    called as decorator with args (<class 'state_chain.Foobar'>,)
    >>> @f(                             #doctest: +ELLIPSIS
    ...     'one long argument',
    ...     'another long argument',
    ... )
    ... def g(): pass
    ...
    called as decorator with args ('one long argument', 'another long argument')
    called as decorator with args (<function g at ...>,)
    >>> @f('one long argument',         #doctest: +ELLIPSIS
    ...    'another long argument')
    ... def g(): pass
    ...
    called as decorator with args ('one long argument', 'another long argument')
    called as decorator with args (<function g at ...>,)
    >>> @f(                             #doctest: +ELLIPSIS
    ...     # A weirdly placed comment
    ...   )
    ... @f
    ... def g(): pass
    ...
    called as decorator with args ()
    called as decorator with args (<function g at ...>,)

    """

    def get_indentation(line):
        for i, c in enumerate(line):
            if not c.isspace():
                break
        return line[:i]

    # First, we try to look at the line where Python says the function call is.
    # Unfortunately, Python doesn't always give us the line we're interested in.
    call_frame = inspect.currentframe().f_back.f_back
    call_info = inspect.getframeinfo(call_frame, context=0)
    source_lines = linecache.getlines(call_info.filename)
    if not source_lines:
        # Reading the source code failed, return the fallback value.
        return default
    try:
        call_line = source_lines[call_info.lineno - 1]
    except IndexError:
        # The source file seems to have been modified.
        return default
    call_line_ls = call_line.lstrip()
    if call_line_ls.startswith('@'):
        # Note: there is a small probability of false positive here, if the
        # function call is on the same line as a decorator call.
        return True
    if call_line_ls.startswith('class ') or call_line_ls.startswith('def '):
        # Note: there is a small probability of false positive here, if the
        # function call is on the same line as a `class` or `def` keyword.
        return True
    # Next, we try to find and examine the line after the function call.
    # If that line doesn't start with a `class` or `def` keyword, then the
    # function isn't being called as a decorator.
    def_lineno = call_info.lineno
    while True:
        try:
            def_line = source_lines[def_lineno]
        except IndexError:
            # We've reached the end of the file.
            return False
        def_line_ls = def_line.lstrip()
        if def_line_ls[:1] in (')', '#', '@', ''):
            def_lineno += 1
            continue
        break
    if not (def_line_ls.startswith('class') or def_line_ls.startswith('def')):
        # Note: there is a small probability of false negative here, as we might
        # be looking at the wrong line.
        return False
    # Finally, we look at the lines above, taking advantage of the fact that a
    # decorator call is at the same level of indentation as the function or
    # class being decorated.
    def_line_indentation = get_indentation(def_line)
    for lineno in range(call_info.lineno - 1, 0, -1):
        line = source_lines[lineno - 1]
        line_indentation = get_indentation(line)
        if line_indentation == def_line_indentation:
            line_ls = line.lstrip()
            if line_ls[:1] in (')', ','):
                continue
            return line_ls.startswith('@')
        elif len(line_indentation) < len(def_line_indentation):
            break
    return default