在运行时更改 **library** 装饰器的参数值（例如 numba 的 ngit(parallel) ）

Question

虽然装饰器接受参数，但当解释器看到底层函数的声明时，它们的值被解析，因此在函数调用期间保持“不变”。

问题是如何在运行时修改装饰器参数值，当您使用来自库 或其他第三方代码的装饰器时。我的问题具体是，如何决定何时允许在 numba 的 njit 装饰器上并行运行，它接受这个作为布尔参数。

---

免责声明

这个问题已经被问了很多，但是每个人都隐含地假设我们控制着装饰器的源代码。在 this question 中，@C2H5OH 的回答实际上解决了这个问题，但是它与“外部”装饰器一起工作的事实几乎没有被注意到，所以任何遇到同样问题的人都必须进行大量搜索。
因此，这个问题旨在自我回答，以查明这种差异并具体说明如何使用 numba 来完成。

---

所以，对于代码：

假设我们必须使用以下装饰器，它在函数调用之前打印其参数的值但我们无法修改其来源：

def decorator_with_argument(a):
    def actual_decorator(func):
        def function_wrapper(*args, **kwargs):
            print(f"Before function call. a={a}")
            func(*args, **kwargs)
        return function_wrapper
    return actual_decorator

我想做这样的事情：

a = 1
@decorator_with_argument(a)
def foo():
    print("Hi")


>>> foo() 
# Before function call. a=1
# Hi

>>> a = 2
>>> foo()
# Before function call. a=1 <---- I would like it to print 2
# Hi

---

对于 numba 案例，

一个小例子如下：

import numba as nb
parallel = True

@nb.njit(parallel=parallel)
def parallel_test(A):
     s = 0
     for i in nb.prange(A.shape[0]):
          s += A[i]
     return s

并且通过更改 parallel 的值，我想 allow/stop 并行执行。

Answer 1

The workaround to the problem is to encapsulate the declaration of the decorated function inside (yet) another wrapper function that handles the "setup". The decorator's argument can be passed as an argument to the outer "setup" function, which will actually redeclare the same functions (but different decorator behavior) each time.

Now there's a scope issue to solve, since our original function is now nested. For that, we can simply return the decorated function or assign it to a global variable.

The code solution:

def function_setup(a):
    @decorator_with_argument(a)
    def _foo():
        print("Hi")

>>> foo = function_setup(1)
>>> foo()
# Before function call. a=1
# Hello

>>> foo = function_setup(2)
>>> foo()
# Before function call. a=2
# Hello

---

For the numba case:

def function_setup(parallel):
    @nb.njit(parallel=parallel)
    def _parallel_test(A):
        s = 0
        for i in nb.prange(A.shape[0]):
            s += A[i]
        return s
    return _parallel_test


import numpy as np
from tqdm import tqdm

A = np.random.random(100)

parallel_test = function_setup(parallel=True)
for _ in tqdm(range(1000000)):
    parallel_test(A)                           
# 100%|██████████| 1000000/1000000 [00:04<00:00, 230644.00it/s]

parallel_test = function_setup(parallel=False)
for _ in tqdm(range(1000000)):
    parallel_test(A)
# 100%|██████████| 1000000/1000000 [00:00<00:00, 1741571.80it/s]

Supposedly, one would like to enable/disable parallel 运行 all njit functions together, so all them will have to be included in the function_setup wrapper and assign to module-level variables.