python decorator learning

decorators.py

import functools
import types

# if passed as 1st argument to decorator will be used for setup and teardown
from typing import Optional, Generator, Dict, Callable, Union


# first call by next will setup. 2nd call will teardown.
def setup():
    print("setup")
    yield {"context_data": "wootus maximus"}
    print("teardown")


# just a normal function
def setup_not_gen():
    print("setup_not_gen")
    return "not a gen"


def pardec_factory(*f_args, **f_kwargs):
    """
    Wraps a function to inject arguments and or add setup and teardown logic.
    If 1st param is a generator it will be removed from injected parameters and used for setup before run and teardown
    after run.
    Positional arguments will be merged with wrapped function arguments before injecting.
    KW arguments wille be merged with wrapped function kw arguments before injecting.
    :return: Same as wrapped function
    """
    print("pardec_factory", f_args, f_kwargs)

    f_args = list(f_args)  # need to convert tuple to list, so we can augment it
    # if first arg is a generator pop it off arg list and assign to gen
    # gen will later be checked if it's a generator for setup / teardown
    if len(f_args) and isinstance(f_args[0], types.GeneratorType):
        gen = f_args.pop(0)
    else:
        gen = None

    # decorator is only ever called with func it is wrapping
    def pardec(func):
        @functools.wraps(func)  # this helps with stack trace
        def wrapper(*p_args, **p_kwargs):  # wrapper gets called with arguments passed to function being decorated
            p_args = list(p_args)  # need to convert tuple to list, so we can augment it
            print("wrapper", p_args, p_kwargs)
            # combine the dictionary args into one, the lowest decorator wins in case of dup keys
            kwargs = p_kwargs.copy()
            kwargs.update(f_kwargs)
            # if gen is a generator pass it kwargs to func with key "init"
            if gen is not None:
                kwargs["init"] = next(gen)

            try:
                # return the value of the func, so we can get results if we want them
                return func(*p_args + f_args, **kwargs)
            finally:
                # if gen was a generator call it again for teardown
                if gen is not None:
                    try:
                        next(gen)
                    except StopIteration:  # only pass on this exception. others should bubble up.
                        pass

        return wrapper

    return pardec


@pardec_factory(setup(), 'x', 'y', 'z', woot="par")
@pardec_factory(setup_not_gen, 'a', 'b', 'c', woot="nope")  # this woot with overwrite the woot above
def myfunc(*args, woot=None, **kwargs):
    print("myfunc", args, kwargs, f"woot={woot}")
    return "all good in the hood"
    # print("myfunc", args, woot)


print("myfunc returned: ", myfunc(1, 2, 3))
print("")


# this decorator can be used with or without parameters
def pardec_factory_opt(_func: Optional[Callable] = None, *, init: Optional[Union[Callable, Generator]] = None,
                       f_kwargs: Optional[Dict] = None):
    """
    Wraps a function to inject arguments and or add setup and teardown logic.
    :param _func: Internal.
    :param init: Optional generator for setup before function run and teardown after function run.
    :param f_kwargs: Optional kwargs that will merge with wrapped function kwargs and be passed to wrapped function.
    :return: Same as wrapped function.
    """
    if f_kwargs is None:  # use default values when not specified
        f_kwargs = {"woot": "default woot"}
    print("pardec_factory", init, f_kwargs)

    if isinstance(init, types.FunctionType):  # if init is a function we have to invoke it to get generator
        gen = init()
    elif isinstance(init, types.GeneratorType):  # init is already a generator, just use it
        gen = init
    else:
        gen = None

    # decorator is only ever called with func it is wrapping
    def pardec(func):
        @functools.wraps(func)  # this helps with stack trace
        def wrapper(*p_args, **p_kwargs):  # wrapper gets called with arguments passed to function being decorated
            print("wrapper", p_args, p_kwargs)
            # combine the dictionary args into one, the lowest decorator wins in case of dup keys
            kwargs = p_kwargs.copy()
            kwargs.update(f_kwargs)
            # if gen is a generator pass it kwargs to func with key "init"
            if gen is not None:
                kwargs["init"] = next(gen)

            try:
                # return the value of the func, so we can get results if we want them
                return func(*p_args, **kwargs)
            finally:
                # if gen was a generator call it again for teardown
                if gen is not None:
                    try:
                        next(gen)
                    except StopIteration:  # only pass on this exception. others should bubble up.
                        pass

        return wrapper

    if _func is None:
        return pardec  # decorator was invoked with parameters
    else:
        return pardec(_func)  # decorator was not invoked in declaration


@pardec_factory_opt  # will use default values
@pardec_factory_opt(init=setup, f_kwargs={"woot2": "par"})  # will do setup and teardown with init generator
@pardec_factory_opt(f_kwargs={"woot2": "nope"})  # this woot2 with overwrite the woot above
def myfunc2(*args, woot=None, **kwargs):
    print("myfunc2", args, kwargs, f"woot={woot}")
    return "all good in the hood"
    # print("myfunc", args, woot)


print("myfunc2 returned: ", myfunc2(1, 2, 3))