用Python装饰整个库

Question

我刚接触装饰器的概念（还在努力理解它们），但我觉得我遇到了一个很适合用装饰器解决的问题。我想要一个类，这个类可以对数学库中的所有函数进行装饰。更具体地说，我的类有两个成员，x 和 flag。当 flag 为真时，我希望调用原始的数学函数。当 flag 为假时，我希望返回 None。

作为我所提问的框架，这就是我的类：

import math

class num(object):
  def __init__(self, x, flag):
    self.x = x
    self.flag = flag

  def __float__(self):
    return float(self.x)

结果，这样做是没问题的：

a = num(3, True)
print math.sqrt(a)

但是在我理想的情况下，这个应该返回 None：

b = num(4, False)
print math.sqrt(b)

有没有什么建议或者技巧可以让这个在整个函数库中都能应用呢？

Answer 1

你可以用装饰器来实现这个功能，不过你不需要使用 @decorator 这种写法。

下面的代码会把你列出的每个函数从 math 模块导入到当前模块的命名空间中，并用你定义的装饰器进行包装。这应该能让你大致明白这个概念。

from functools import wraps
def check_flag(func):
    @wraps(func)
    def _exec(x, *args, **kw):
        if getattr(x, 'flag', False):
            return None

        return func(x, *args, **kw)

    return _exec

import sys, math
_module = sys.modules[__name__]
for func in ('exp', 'log', 'sqrt'):
    setattr(_module, func, check_flag(getattr(math, func)))

你 可以 像 Alex 演示的那样，自动列出 math 模块中定义的函数，但我觉得明确地只包装你想用的那些函数会更好。

Answer 2

大致的意思是这样的……：

>>> class num(object):
...   def __init__(self, x, flag):
...     self.x = x
...     self.flag = flag
...   def __float__(self):
...     return float(self.x)
...   from functools import wraps
>>> def wrapper(f):
...   @wraps(f)
...   def wrapped(*a):
...     if not all(getattr(x, 'flag', True) for x in a):
...       return None
...     return f(*(getattr(x, 'x', x) for x in a))
...   return wrapped
... 
>>> import inspect
>>> import math
>>> for n, v in inspect.getmembers(math, inspect.isroutine):
...   setattr(math, n, wrapper(v))
... 

>>> a = num(3, True)
>>> print math.sqrt(a)
1.73205080757
>>> b = num(4, False)
>>> print math.sqrt(b)
None

需要注意的是，这个包装器也适用于math模块中的非一元函数（如果任何一个参数的.flag为False，就会返回None），并且允许混合调用（也就是说，有些参数可以是num的实例，其他的可以是真正的浮点数）。

关键的部分是，适用于任何“将某个模块中的所有函数都包装起来”的任务，就是使用inspect模块来获取math模块中所有函数的名称和值（无论是内置的还是自定义的），然后明确调用这个包装器（和装饰器的语法效果一样），将这个名称设置为math模块中被包装的值。