使用Python装饰器进行代码重构？

Question

我现在在处理一段代码时遇到了一些困难。我知道这段代码可以进行重构，但我找不到一个既好又聪明的优雅解决方案。

这里有两个函数（我的代码中还有很多类似的函数）：

def fooA(param1, param2):
    if param2 == True:
       code_chunk_1

    fooA_code  #uses only param1

    if param2 == True:
       code_chunk_2


def fooB(param1, param2):
    if param2 == True:
       code_chunk_1

    fooB_code  #uses only param1

    if param2 == True:
       code_chunk_2

我最初的想法是使用这个装饰器：

def refactorMe(func):
    def wrapper(*args):
        if args[-1]:
            code_chunk_1

        func(*args)

        if args[-1]:
            code_chunk_2

    return wrapper

最后：

@refactorMe
def fooA(param1, param2):
    fooA_code  #uses only param1

@refactorMe
def fooB(param1, param2):
    fooB_code  #uses only param1

不幸的是，我对这个解决方案并不满意：

• 这个装饰器是“侵入式”的，而且只适用于fooA和fooB这两个函数
• param2在fooA和fooB的主体中不再使用，但我们必须在函数签名中保留它

也许我没有按照装饰器最初的目的来使用它？

有没有其他方法可以重构这段代码？

非常感谢！

Answer 1

因为你想在传入的选项为真时启用一些包装功能，所以可以考虑使用关键字参数。下面是一个实际的例子，如果需要的话，它会把你的代码放在一个（数据库）事务中：

def wrap_transaction(func):
    def wrapper(*args, **kwargs):
        # If the option "use_transaction" is given, wrap the function in
        # a transaction.  Note that pop() will remove the parameter so
        # that it won't get passed to the wrapped function, that does not need
        # to know about its existance.
        use_transaction = kwargs.pop('use_transaction', False)

        if use_transaction:
            get_connection().begin_transaction()

        try:
            result = func(*args, **kwargs)
        except:
            if use_transaction:
                get_connection().rollback()
            raise

        if use_transaction:
            get_connection().commit()

        return result

    return wrapper

@wrap_transaction
def my_func(param):
    # Note that this function knows nothing about the 'use_transaction' parameter
    get_connection().exec("...")


# Usage: Explicitely enabling the transaction.
my_func(param, use_transaction=True)

Answer 2

这样怎么样：

def call_one(func, param1, param2):
    if param2:
        code_chunk_1

    func(param1)

    if param2:
        code_chunk_2

def _fooA(param1):
    fooA_code  #uses only param1

def _fooB(param1):
    fooB_code  #uses only param1

def fooA(param1, param2):
    call_one(_fooA, param1, param2)

def fooB(param1, param2):
    call_one(_fooB, param1, param2)

Answer 3

你描述的情况是有一些重复的代码和行为，然后又是一些重复的代码。简单来说，这种情况可以用一个叫做高阶函数的东西来处理，比如 map、reduce 或 filter。

你可以按照 Ned 的建议去做（不过，我会用functools.partial，而不是长篇大论地定义 fooA/fooB）：

import functools

...

fooA = functools.partial(call_one, _fooA)
fooB = functools.partial(call_one, _fooB)

... 但这样做实际上又回到了你用装饰器的同样情况，同时还在命名空间里增加了一些杂乱的东西。

你可以重写你的装饰器，让它只接受一个参数的函数，但返回的函数却可以接受两个参数：

def refactorMe(func):
    def wrapper(parm1, parm2):
        if parm1:
            code_chunk_1

        func(parm1)

        if parm2[-1]:
            code_chunk_2

    return wrapper

去掉那些复杂的星号操作是个改进，因为这个装饰器并不适用于所有函数，所以我们应该明确这一点。我喜欢我们减少参数数量的做法，因为任何查看代码的人可能会被调用函数时多加的一个参数搞糊涂。此外，感觉上改变被装饰函数的参数签名的装饰器应该是不太好的做法。

总结一下：

装饰器是高阶函数，而模板化行为正是它们的用途。

我会接受这个代码是专门针对你的 fooXXX 函数的，做一个内部的装饰器，并让它只接受所需的参数数量（因为 foo(*args, **kwargs) 的签名会让人很难理解）。

def _refactorMe(func):
        @functools.wraps(func) #the wraps decorator propagates name/docsting
        def wrapper(parm1, parm2):
            if parm1:
                code_chunk_1

            func(parm1, parm2)

            if parm2:
                code_chunk_2

        return wrapper

我会保留两个参数的调用，即使其中一个没有用，这样装饰器就不会改变函数的签名。虽然这并不是绝对必要的，因为如果你在文档中说明了函数在装饰后的样子，并且只将装饰器限制在这小部分函数上，那么签名的变化就不那么重要了。

@_refactorMe
def fooB(param1, param2):
    fooB_code  #uses only param1


@_refactorMe
def fooB(param1, param2):
    fooB_code  #uses only param1