Home > Enterprise >  Multiple decorators on single function
Multiple decorators on single function

Time:09-15

How can I add multiple decorators to a function so that each configuration will combine into one final dictionary?

My explanation probably doesn't make sense, but look at the example.


def instruction(name):
  def decorator(callback):
    def wrapper(self, *args, **kwargs):
      # somehow combine every additional dict into this dict
      return {
        "name": name,
        "callback": lambda: callback(self, *args, **kwargs)
      }
    return wrapper
  return decorator

def layout(config):
  def decorator(_):
    return {"layout": config}
  return decorator

# example decorator
def documentation(text):
  def decorator(_):
    return {"documentation": text}
  return decorator

@instruction("poke")
@layout({"address": 2, "data": 1})
@documentation("handle poke instruction") # additional decorators
def instruction_poke(self):
  address = self.serial_read_bytes(2)
  data = self.serial_read_bytes(1)
  print(f"poke {address=} {data=}")
  # ...

def main():
  self = ... # object with the serial_read_bytes method
  instruction = instruction_poke(self)
  print(instruction)
  # expected output
  # {
  #   "name": poke,
  #   "documentation": "handle poke instruction",
  #   "layout": {"address": 2, "data": 1},
  #   "callback": <lambda to original function>
  # }

if __name__ == "__main__":
  main()

I'm having trouble with this as multiple decorators has a stack-like behaviour.

My current approach is just to store every decorator's fields in a constants.py file, which in my opinion, is troublesome to mess around with several files.

I appreciate your time :)

CodePudding user response:

You can create another level of closure to generalize the creation of decorator configurer functions.

The wrapper function should distinguish between an actual callback function and a stacked decorator, which can be done by testing if the code object of the callback function is that of the wrapper function itself, in which case it should simply return the dict returned by callback function with the given key-value pair added:

def add_config(key):
    def configurer(value):
        def decorator(callback):
            def wrapper(*args, **kwargs):
                if callback.__code__ is wrapper.__code__:
                    return {key: value, **callback(*args, **kwargs)}
                return {key: value, 'callback': lambda: callback(*args, **kwargs)}
            return wrapper
        return decorator
    return configurer

instruction, layout, documentation = map(
    add_config, ['name', 'layout', 'documentation']
)

@instruction('poke')
@layout({"address": 2, "data": 1})
@documentation("handle poke instruction")
def instruction_poke(i):
    return i   1

instruction = instruction_poke(2)
print(instruction)
print(instruction['callback']())

This outputs:

{'name': 'poke', 'layout': {'address': 2, 'data': 1}, 'documentation': 'handle poke instruction', 'callback': <function add_config.<locals>.configurer.<locals>.decorator.<locals>.wrapper.<locals>.<lambda> at 0x7fd9476fad30>}
3

Demo: https://replit.com/@blhsing/SadHugeEquation

CodePudding user response:

Depending on how you want to go about this, there are multiple options available. The one by @blhsing certainly works well. Here are some more:

Wrap in dictionary (simple)

def wrap_with_info(**kwargs):
    def decorator(obj):
        if isinstance(obj, dict):
            return obj | kwargs
        # `obj` must be a function
        assert callable(obj)
        return kwargs | {"callback": obj}
    return decorator

This allows you to decorate like this:

@wrap_with_info(
    name="poke",
    documentation="handle poke instruction",
    layout={"address": 2, "data": 1},
)
def your_function():
    pass

And also like this:

@wrap_with_info(name="poke")
@wrap_with_info(documentation="handle poke instruction")
@wrap_with_info(layout={"address": 2, "data": 1})
def your_function():
    pass

Note that this method always essentially turns your_function into a dictionary:

print(your_function)

Gives this output:

{'layout': {'address': 2, 'data': 1}, 'callback': <function your_function at 0x7fd1184ae710>, 'documentation': 'handle poke instruction', 'name': 'poke'}

Decorate to return dictionary

from functools import wraps


def add_info(**kwargs):
    def decorator(func):
        def wrapper(*func_args, **func_kwargs):
            @wraps(func)
            def callback():
                return func(*func_args, **func_kwargs)
            return kwargs | {"callback": callback}
        return wrapper
    return decorator

Then you also decorate like this:

@add_info(
    name="poke",
    documentation="handle poke instruction",
    layout={"address": 2, "data": 1},
)
def another_function():
    pass

But now another_function remains a callable that returns the dictionary containing the original function (wrapped) and the additional data:

print(another_function())

Gives this output:

{'name': 'poke', 'documentation': 'handle poke instruction', 'layout': {'address': 2, 'data': 1}, 'callback': <function another_function at 0x7fb4ae7a67a0>}

The functools.wraps decorator ensures that the original function's signature is retained when we stick it in the callback. This is why you see its name in the output instead of something like ...local.callback....


(Side note: I am using the | notation to combine dictionaries, which is available since Python 3.9. If that doesn't work for you, you can use other means, like dict.update for example.)


Advantages

It seems like all you are doing with your decorators is adding key-value-pairs to a dictionary of information around an actual function. This method makes it unnecessary to use multiple decorators. Instead you can pass any number of key-value-pairs depending on the function.

They can also be constructed dynamically, if that is something you need, and if you know the data you want to add to the functions ahead of time:

info_x = {"spam": 1, "eggs": 2}
info_y = {"info": "abc"}
info_z = {"something": ["else"]}

...

@add_info(**info_x, **info_y)
def foo():
    pass


@add_info(**info_y, **info_z)
def bar():
    pass


@wrap_with_info(**info_z, **info_x)
def baz():
    pass


print(foo())
print(bar())
print(baz)

Gives this output:

{'spam': 1, 'eggs': 2, 'info': 'abc', 'callback': <function foo at 0x7f5d1abb29e0>}
{'info': 'abc', 'something': ['else'], 'callback': <function bar at 0x7f5d1abb29e0>}
{'something': ['else'], 'spam': 1, 'eggs': 2, 'callback': <function baz at 0x7efd492aeb00>}
  • Related