metadecorators to construct decorators that only transform function arguments or return values

transform_decorator.py

from functools import wraps, partial

def transform_return(transformation, decorated=None):
    if decorated is None:
        return partial(transform_return, transformation)
    @wraps(decorated)
    def wrapper(*args, **kwargs):
        return transformation(decorated(*args, **kwargs))
    return wrapper

def transform_arguments(transformation, decorated=None):
    if decorated is None:
        return partial(transform_arguments, transformation)
    @wraps(decorated)
    def wrapper(*args, **kwargs):
        args2, kwargs2 = transformation(*args, **kwargs)
        return decorated(*args2, **kwargs2)
    return wrapper

if __name__ == '__main__':

    @transform_return(''.join)
    def some_string_builder(seq):
        for i, item in enumerate(seq, start=1):
            if i > 1:
                yield '; '
            yield 'Item '
            yield str(i)
            yield ': '
            yield str(item)


    # alternatively, use it to make a reusable decorator:
    @transform_arguments
    def sum_args(*args):
        return args, {'sum': sum(args)}

    import random
    from collections import Counter

    @sum_args
    def pick(*args, sum):
        random_index = random.randrange(sum)
        cum_sum = 0
        for i, item in enumerate(args):
            cum_sum += item
            if random_index < cum_sum:
                return i

    print(some_string_builder([1, 22, 333, 4444, 55555]))
    print(Counter(pick(100, 1, 0, 100) for _ in range(100000))) # should print 0 or 3 almost 50% of the time, 1 almost 0.5% of the time, and never 2


# transform_arguments :: Callable[
#     [Callable[ArgSpec1, ArgSpec2]],
#      Callable[
#         [Callable[ArgSpec1, Out]],
#          Callable[ArgSpec2, Out]]]


# transform_return :: Callable[
#     [Callable[[In], Out]],
#      Callable[
#         [Callable[ArgSpec, In]],
#          Callable[ArgSpec, Out]]]