Creating very similar functions with different parameters

[Steven D'Aprano]

On Mon, 24 Oct 2011 16:29:25 -0500, Andrew Berg wrote:

> 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




-- 
Steven