使用Python装饰器启用/禁用程序的某些部分

时间:2018-09-07 12:38:14

标签: python decorator

我有一堆脚本,这些脚本本质上是数据准备步骤,用于设置仿真模型的数据。我通常只想运行其中的一部分,例如“ phase1”或“ phase2”,但是大多数“ phases”不止一行,因此注释掉不是很方便。所以我通常会这样做:

# Phase 1
if True:
  do_step_1('high')
  do_step_2()
  for i in range(1,10):
    do_step_3()

#Phase 2
if True:
  do_step_1('low')
  do_something_else()

然后根据需要将True更改为False。

现在,这很麻烦。有时,阶段彼此依赖(因此当我运行3时,我也需要运行1),它们是嵌套的,等等。

我想做的是,有某种方式可以将参数传递给脚本,该参数将运行一个或多个“阶段”,而我需要某种方式来“标记”某些功能,代码块或作用域作为一部分这个“阶段”。某个代码块可以是多个阶段的一部分,因此当存在依赖于块A的块B和C时,我可以将A标记为“ phase1”和“ phase2”的一部分,然后在运行phase1时将其标记为将会运行A块和B块,对于阶段2,将运行A块和C块。我希望这仍然有意义。

因此,我一直认为装饰者可以完美地做到这一点,这样我就可以(在概念上)做类似的事情

@partOfAPhase("phase1", "phase2")
def f1():
    pass

然后以某种方式,我将要运行的“阶段”列表传递给我的程序(从命令行,或通过将其设置为某个配置变量),并且在我的程序运行时,它仅执行那些被装饰为指定要运行的阶段之一。

所以,我认为我需要的是一个通用修饰符,该修饰符可以应用于接受任意数量参数的函数或成员函数,并且我需要能够传递'标签”到装饰器本身。然后在装饰器内部,我需要检查(在调用原始函数或成员时)该装饰器的标签是否存在于要运行的全局标签列表中(也许是静态类?)。

我看着https://gist.github.com/Zearin/2f40b7b9cfc51132851a,似乎最后还是或多或少地做了我想做的事,但是我不能完全困惑所有事情一起做我想做的事。更具体地说,我不太了解双重嵌套的装饰器生成器,我是否需要两个函数还是仅一个函数来实现此功能,以及如何访问传递给装饰器的参数(即要运行的阶段) )。

2 个答案:

答案 0 :(得分:1)

不确定是否满足您的需求,但这是一个快速而肮脏的概念证明:

# This first part could go in its own module or a class or whatever
_program = []

def partOfPhase(*phases):
    def decorator(fn):
        _program.append((fn, tuple(phases)))
        return fn
    return decorator

def partOfProgram(fn):
    _program.append((fn, None))
    return fn

def runProgram(phase):
    for fn, fn_phases in _program:
        if fn_phases is None or any(p in fn_phases for p in phases):
            fn()

# This is the actual script
import sys

@partOfPhase('phase1')
def step1():
    print('step1')

@partOfPhase('phase1', 'phase2')
def step2():
    print('step2')

@partOfProgram
def step3():
    print('step3')

@partOfPhase('phase2')
def step4():
    print('step4')

if __name__ == '__main__':
    phases = sys.argv[1:]
    runProgram(phases)

例如,如果将其另存为phases.py,则会得到:

> python phases.py
step3

> python phases.py phase1
step1
step2
step3

> python phases.py phase2
step3
step4

> python phases.py phase1 phase2
step1
step2
step3
step4

编辑

我认为这可能更像您在想什么,根据阶段而禁用的功能:

# This first part could go in its own module or a class or whatever
from functools import wraps

_enabledPhases = []

def enablePhase(*phases):
    _enabledPhases.extend(phases)

def partOfPhase(*phases):
    def decorator(fn):
        @wraps(fn)  # Just "cosmetic" wrapping
        def decorated(*args, **kwargs):
            if any(p in phases for p in _enabledPhases):
                fn(*args, **kwargs)
        return decorated
    return decorator

# This is the actual script
import sys

@partOfPhase('phase1')
def step1():
    print('step1')

@partOfPhase('phase1', 'phase2')
def step2():
    print('step2')

def step3():
    print('step3')

@partOfPhase('phase2')
def step4():
    print('step4')

if __name__ == '__main__':
    phases = sys.argv[1:]
    enablePhase(*phases)
    step1()
    step2()
    step3()
    step4()

答案 1 :(得分:1)

在找到链接中显示的方法非常复杂并且使用可调用装饰器对象执行此操作之后,这就是我的最终结论。这适用于带有任意数量参数的独立函数和方法(带有一些简单的测试,这些测试仅打印易于视觉验证的结果):

import functools
import sys

class runConditional(object):
    def __init__(self, datasets):
        self.datasets = datasets

    def __call__(self, func):
        def wrapped_f(*args, **kwargs):
            global to_run
            for d in self.datasets:
                if d in to_run:
                    sys.stdout.write(" 1")
                    func(*args, **kwargs)
                    return
            sys.stdout.write(" 0")

        return wrapped_f

@runConditional([1])
def fun1():
    pass

@runConditional([2])
def fun2():
    pass

@runConditional([1,2,3])
def fun3(arg1, arg2):
    pass

def fun_always():
    sys.stdout.write(" 1")
    pass

@runConditional([])
def fun_never():
    pass

class test():
    @runConditional([1])
    def m1(self):
        pass

    @runConditional([2])
    def m2(self):
        pass

    @runConditional([1,2,3])
    def m3(self, arg1):
        pass

    def m_always(self):
        sys.stdout.write(" 1")
        pass

    @runConditional([])
    def m_never(self):
        pass

def run_test(funcs_to_run, expected):
    global to_run
    t = test()
    funcs = [ fun1, fun2, functools.partial(fun3, "123", "meh"), fun_always, fun_never,
            t.m1, t.m2, functools.partial(t.m3, "321"), t.m_always, t.m_never ]
    to_run = funcs_to_run
    print "Expected: " + " ".join(map(str, expected))
    sys.stdout.write("Actual:  ")
    for f in funcs:
        f()
    print ""
    print ""

run_test([2],       [ 0, 1, 1, 1, 0,  0, 1, 1, 1, 0 ])
run_test([1],       [ 1, 0, 1, 1, 0,  1, 0, 1, 1, 0 ])
run_test([],        [ 0, 0, 0, 1, 0,  0, 0, 0, 1, 0 ])
run_test([1, 2],    [ 1, 1, 1, 1, 0,  1, 1, 1, 1, 0 ])