Code refactoring with python decorators?

Question

Hi everyone,

I'm actually struggling with some piece of code. I do know that it can be refactored, but I can't find the nice-smart-elegant solution.

Here are two functions (much more functions of that kind are in my code):

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

My first idea was to use this decorator:

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

        func(*args)

        if args[-1]:
            code_chunk_2

    return wrapper

And finally:

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

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

Unfortunately, I'm not happy with this solution:

• This decorator is "intrusive" and specific to the fooA & fooB functions
• param2 is not used anymore in the fooA & fooB body, but we must keep it in the function signature

Perhaps I'm not using the decorator for its initial purpose?

Is there any other way to refactor the code?

Thanks a lot!

Answer 1

How about:

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 2

Since you try to enable some wrapper functionality iff the passed option is True, consider using keyword arguments. Here is a real-world example that will wrap your code in a (database-) transaction if requested:

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 3

I would do a straightforward extract method refactoring:

def _code_chunk_1(param):
    if param == True:
        code_chunk_1

def _code_chunk_2(param):
    if param == True:
        code_chunk_2

def fooA(param1, param2):
    _code_chunk_1(param2)

    fooA_code  #uses only param1

    _code_chunk_2(param2)

def fooB(param1, param2):
    _code_chunk_1(param2)

    fooB_code  #uses only param1

    _code_chunk_2(param2)

The decorator looks inappropriate to me in this context. Ned's answer above also looks nice.

Answer 4

I'm wondering if you're adding debugging code. Since param2 is not used in the function proper, maybe you want to move it into the decorator:

class debugging:
    def __init__(self, show):
        self.show = show

    def __call__(self, f):
        def wrapper(*args):
            if self.show:
                print "inside", f

            rv = f(*args)

            if self.show:
                print "outside", f

            return rv

        return wrapper

@debugging(True)
def test(n):
    print n

test(10)

will print

inside <function test at 0x7fb28ff102a8>
10
outside <function test at 0x7fb28ff102a8>

Answer 5

What you are describing is a situation where you have some boilerplate, some behaviour, followed by some boiler plate. Essentially a situation where you could use a Higher Order Function (like map, reduce or filter).

You could do what Ned suggests (though, I'd use functools.partial rather than defining fooA/fooB longhand):

import functools

...

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

... but that effectively gets you back to the same place as with your decorator, introducing some clutter into the namespace along the way.

You could rewrite your decorator to allow functions that only take one parameter, but return functions that take two:

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

        func(parm1)

        if parm2[-1]:
            code_chunk_2

    return wrapper

Getting rid of the star magic is an improvement as this decorator is not general to all functions so we should be explicit about it. I like the fact that we change the number of parameters less as anyone looking at the code could easily be confused by the fact that when we call the function we are adding an extra parameter. Furthermore it just feels like decorators that change the signature of the function they decorate should be bad form.

In summary:

Decorators are higher order functions, and templating behaviour is precisely what they're for.

I would embrace the fact that this code is specific to your fooXXX functions, by making the decorator internal and having it take precisely the number of arguments needed (because foo(*args, **kwargs) signatures makes introspection a pain).

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

I'd leave the calls taking two parameters, even though one is unused just so that the decorator doesn't change the signature. This isn't strictly necessary as if you document the functions as they look after decoration and you are restricting the use of the decorator to this small set of functions then the fact that the signature changes shouldn't be that big a deal.

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


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

Answer 6

I like Ferdinand Beyer's answer and I think we need examples like that to understand what we are talking about. I'm just goin to give two further inspirational suggestions.

Why not explicitly use transaction code?

def fooA(param1, use_transaction=param2):
    enter_transaction(param2)
    fooA_code  #uses only param1
    exit_transaction(param2)

def fooB(param1, use_transaction=param2):
    enter_transaction(param2)
    fooB_code  #uses only param1
    exit_transaction(param2)

Now with that written we understand that we should probably write this:

def fooA(param1, use_transaction=param2):
    with transaction(param2):
        fooA_code  #uses only param1

def fooB(param1, use_transaction=param2):
    with transaction(param2):
        fooB_code  #uses only param1

Using some context manager.

But wait! We can put that outside!

if you want this use:

with transactional():
    fooA(param1)

for the not param2 case, simply call fooA(param1)

Last syntax suggestion, when param2 == true:

do_transaction(fooA, param1)

here we define

def do_transaction(func, *args):
    code_1
    func(*args)
    code_2

Ok that was my stream of thoughts. Can you use a context manager? It is also hard to document, but somehow this wrapping process must be integral to your application, or if it's not, you could remove it.