Creating very similar functions with different parameters

[Andrew Berg]

I want to create a decorator with two different (but very similar)
versions of the wrapper function, but without copying giant chunks of
similar code. The main difference is that one version takes extra
parameters.

def test_dec(func, extra=False):
if extra:
def wrapper(ex_param1, ex_param2, *args, **kwargs):
print('bla bla')
print('more bla')
print(ex_param1)
print(ex_param2)
func(ex_param1, ex_param2, *args, **kwargs)
else:
def wrapper(*args, **kwargs):
print('bla bla')
print('more bla')
func(*args, **kwargs)
return wrapper

If the function I'm wrapping takes ex_param1 and ex_param2 as
parameters, then the decorator should print them and then execute the
function, otherwise just execute the function. I'd rather not write
separate wrappers that are mostly the same.

 

[alex23]

Guido wrote an article on a quick and easy multimethod implementation
using a decorator:
http://www.artima.com/weblogs/viewpost.jsp?thread=101605

 

[Steven D'Aprano]

I want to create a decorator with two different (but very similar)
versions of the wrapper function, but without copying giant chunks of
similar code. The main difference is that one version takes extra
parameters.

def test_dec(func, extra=False):
if extra:
def wrapper(ex_param1, ex_param2, *args, **kwargs):
print('bla bla')
print('more bla')
print(ex_param1)
print(ex_param2)
func(ex_param1, ex_param2, *args, **kwargs)
else:
def wrapper(*args, **kwargs):
print('bla bla')
print('more bla')
func(*args, **kwargs)
return wrapper

If the function I'm wrapping takes ex_param1 and ex_param2 as
parameters, then the decorator should print them and then execute the
function, otherwise just execute the function. I'd rather not write
separate wrappers that are mostly the same.

In principle you could inspect the function's calling signature using the
inspect module, and then decide how to decorate it from that. But I
recommend against that: it's too much like magic.

This is how I would do it:


from functools import wraps

def decorator_factory(extras=None):
"""Factory function which returns a decorator.

Usage:

@decorator_factory(): # Note you need the parentheses.
def spam(...):
return ...

@decorator_factory((extra1, extra2)):
def spam(...):
return ...

"""
# Inner function performing common functionality.
def preamble():
print('bla bla')
print('more bla')
# Decide what sort of decorator is needed.
if extras is None:
# Create decorator style 1.
def decorator(func):
@wraps(func)
def inner(*args, **kwargs):
preamble()
return func(*args, **kwargs)
return inner
else:
# Create decorator style 2.
extra1, extra2 = extras
def decorator(func):
@wraps(func)
def inner(*args, **kwargs):
preamble()
print(extra1)
print(extra2)
return func(extra1, extra2, *args, **kwargs)
return inner
return decorator



If you don't like nested functions inside nested functions, you can pull
out the inner functions:


def preamble():
print('bla bla')
print('more bla')

def plain_decorator(func):
@wraps(func)
def inner(*args, **kwargs):
preamble()
return func(*args, **kwargs)
return inner

def extra_decorator(func):
@wraps(extra1, extra2, func)
def inner(*args, **kwargs):
preamble()
print(extra1)
print(extra2)
return func(extra1, extra2, *args, **kwargs)
return inner


def decorator_factory(plain):
if plain:
return plain_decorator
else:
return extra_decorator