具有函數參數的Python裝飾器

Question

我有一個簡單的函數，定義如下：

def simple_function(x):
    """ x is an input numpy array"""
    return x + 0.1

我想通過對其應用一些邊界條件來修改此函數。 這些邊界條件本身就是x的函數：

def upper_bound(x):
    return x**2

def lower_bound(x):
    return np.zeros(len(x))

特別是，如果simple_function(x)超過upper_bound(x)的值，或低於lower_bound(x) ，我希望裝飾版的simple_function(x)返回值upper_bound(x) ，同樣地，返回lower_bound。 如何使用python中的@decorator語法完成此行為？

Answer 1

如果參數，范圍和結果都是numpy數組，則可以執行幾個數組分配，以將每個元素限制在您的upper_bound和lower_bound函數返回的對應值之間。 核心部分是：

r = f(x)
l = lower_bound(x)
u = upper_bound(x)

i = r < l
j = r > u

r[i] = l
r[j] = u

i和j將是布爾數組，表示需要分別將哪些索引限制在上下邊界。 要使此代碼作為裝飾器工作，只需將其放在一對嵌套函數中，如下所示：

def clamp(f):
    @functools.wraps(f)
    def wrapper(x):
        r = f(x)
        l = lower_bound(x)
        u = upper_bound(x)

        i = r < l
        j = r > u

        r[i] = l
        r[j] = u

        return r

    return wrapper

functools.wraps可以使包裝函數復制修飾后的函數的名稱，注釋和文檔字符串。

上面的代碼假定您始終使用相同的upper_bound和lower_bound函數。 如果需要針對要裝飾的不同功能進行自定義，則可以添加額外的嵌套層，並像Ignacio Vazquez-Abrams的答案中那樣定義“裝飾器工廠”：

def clamp(lower_bound, upper_bound):   # this is the decorator factory function
    def decotator(f):                  # this is the decorator function
        @functools.wraps(f)
        def wrapper(x):                # this is the wrapper function
            ... # same code here as above
            return r

        return wrapper

    return decorator

Answer 2

除了修改__doc__ ，您還可以：

def constrain(lower, upper):
  def outer(f):
    def inner(x):
      r = f(x)
      u = upper(x)
      if r > u:
        return u
      l = lower(x)
      if r < l:
        return l
      return r
    return inner
  return outer

...

@constrain(lower_bound, upper_bound)
def simple_function(x):
   ...

不會處理不同的類型以及下限高於上限的下限。