I wanted to overwrite wraps from pint.UnitRegistry in order to invoke it only if the user used arguments with units. My attempt is here:

from pint import UnitRegistry, Quantity

ureg = UnitRegistry()


def units(*ureg_args):
    def decorator(func):
        @ureg.wraps(*ureg_args, strict=False)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)

        def final_wrapper(*args, **kwargs):
            if any([isinstance(arg, Quantity) for arg in args]) or any(
                [isinstance(arg, Quantity) for arg in kwargs.values()]
            ):
                return wrapper(*args, **kwargs)
            else:
                return func(*args, **kwargs)

        return final_wrapper

    return decorator


@units(“second”, “second”)
def get_time(time):
    return 2 * time

However, this raises the error TypeError: wrapper takes 2 parameters, but 1 units were passed. I guess it’s because ureg.wraps inspects the input arguments, which in this case would be *args, **kwargs and not time. Is there a good workaround for this?