How to to structure python decorating functions to be more maintainable?

Question

I am really tired of plumbing values through decorating functions in my code, and while I think there might be python solutions for me using splat (**) or writing a decorator, I'm not sure what the best type-checker-friendly solution is. I'm feeling frustrated by maintaining code like

@dataclass
class FooBar:
    foo: FooOut
    bar: BarOut

def foo(arg1: str, arg2: int) -> FooOut:
    do_something()
    return inner_func(arg1=arg1, arg2=arg2).foo

def bar(arg1: str, arg2: int) -> BarOut:
    bar = inner_func(arg1=arg1, arg2=arg2).bar
    bar2 = do_some_mutation(bar, arg1)
    return bar2

def inner_func(arg1: str, arg2: int) -> FooBar:
    # do something with the arguments
    return FooBar(foo=some_value, bar=some_other_value)

Now the trouble comes in when inner_func is going to need an annoying new arg3: Optional[str] parameter. The last thing I want to do is add arg3: Optional[str] = None everywhere and then pass arg3 around manually.

One solution is to turn the dictionary of arguments into an object. This is OK except that it will require modifying all of the call sites. A kinda messy variant of this would be to create new functions foo_new(in: InputParams) which take in the input object and have the old functions delegate to the new ones. Then we'd deprecate the old functions eventually.

Another solution is just convert the signature to each of the functions to use **kwargs. This would work except for keeping the static type checker happy. In order for someone to know what valid params there are it would require digging into inner_func to see what is used.

I think the ideal solution would be if I could do something like

@dataclass
class InputParams:
    arg1: str
    arg2: int
    arg3: Optional[str]

def bar(input: **InputParams) -> BarOut:
    bar = inner_func(input)
    bar2 = do_some_mutation(bar, input.arg1)
    return bar

What would you do? Thank you very much!

Answer 1

It seems like you have almost answered your question already, all your functions appear to be something that acts as a light wrapper (with pre- or post-processing only) of inner_func. The fact that you use arg1 in bar is a little sad, and makes things more complicated than they could be (e.g. you can't ignore the arguments).

Attempts to deal with this generically spun a bit out of control, but should be reusable at least (hopefully a canned solution for this can replace blocks of my code). Looked at wrapper factors, generators, and classes, ended up with this.

Example usage:

def inner(a, b=2, c=3):
    return a * b + c


@preprocessor(inner)
def double_a(arguments):
    arguments.arguments['a'] *= 2

@postprocessor(inner)
def add_b_to_result(res, arguments):
    return res + arguments.arguments['b']


def clean_double_a(arguments):
    arguments.arguments['a'] *= 2


def clean_add_b_to_result(res, arguments):
    return res + arguments.arguments['b']


both = processor(inner, pre=clean_double_a, post=clean_add_b_to_result)

@preprocessor(inner)
def both_alt(arguments):
    return clean_double_a(arguments)

@both_alt.postprocessor
def both_alt_post(res, arguments):
    return clean_add_b_to_result(res, arguments)


def self_tests():
    assert inner(1)           == 1 * 2 + 3
    assert double_a(1)        == 2 * 2 + 3
    assert add_b_to_result(1) == 1 * 2 + 3 + 2
    assert both(1)            == 2 * 2 + 3 + 2
    assert both_alt(1)        == 2 * 2 + 3 + 2
    print('Tests complete')

self_tests()  # Tests complete

Helper library:

from functools import partial, update_wrapper
from inspect import signature


class Processor(object):
    def __init__(self, func, pre=None, post=None):
        self.__func = func
        self.__sig = signature(func)
        self.__preprocessor = pre
        self.__postprocessor = post
        if pre or post:
            # Make this instance look like the first of pre/post
            self.__wrapped_by_pre_post = True
            update_wrapper(self, pre or post)
        else:
            # Make this instance look like the function
            self.__wrapped_by_pre_post = False
            update_wrapper(self, func)

    def __call__(self, *args, **kwargs):
        arguments = self.__sig.bind(*args, **kwargs)
        arguments.apply_defaults()
        if self.__preprocessor is not None:
            new_args = self.__preprocessor(arguments)
            if new_args is not None:
                arguments = new_args
        res = self.__func(*arguments.args, **arguments.kwargs)
        if self.__postprocessor is not None:
            res = self.__postprocessor(res, arguments)
        return res

    def preprocessor(self, func):
        if self.__preprocessor is not None:
            raise RuntimeError('defined multiple preprocessors')
        self.__preprocessor = func
        self.__wrap_if_first(func)
        return self

    def postprocessor(self, func):
        if self.__postprocessor is not None:
            raise RuntimeError('defined multiple postprocessors')
        self.__postprocessor = func
        self.__wrap_if_first(func)
        return self

    def __wrap_if_first(self, func):
        # Make it look like the first to process, like getter/setters
        if not self.__wrapped_by_pre_post:
            update_wrapper(self, func)
            self.__wrapped_by_pre_post = True


processor = Processor


def preprocessor(func):
    return Processor(func).preprocessor


def postprocessor(func):
    return Processor(func).postprocessor