Managing attributes inside functions with decorators

user_2713733 0
Viewed 41

This is a toy example of a larger code I'm working right now. Suppose I have a function and I want to decorate it in a way that I can keep record of arguments that are valid, legacy or forbidden:

@option("1")
@option("2", legacy=True)
def do_stuff(opt):
    print(f"You chose {opt}!")


do_stuff("1")
do_stuff("2")
do_stuff("3")

The result of this would be something like

You chose 1!
2 is deprecated!
You chose 2!
3 is not allowed!

I tried implementing this behavior in some different ways now, my latest attempt is this:

from functools import wraps


def option(opt, legacy=False):
    def add_opt(f, option):
        if not hasattr(f, "opts"):
            f.opts = []
        f.opts.append(option)

    def add_legacy_opt(f, option):
        if not hasattr(f, "legacy"):
            f.legacy = []
        f.legacy.append(option)

    def decorate(func):
        @wraps(func)
        def wrapper(option):
            if hasattr(wrapper, "opts") and option not in wrapper.opts:
                print(f"{option} is not allowed!")
                return
            if hasattr(wrapper, "legacy") and option in wrapper.legacy:
                print(f"{option} is deprecated!")
            return func(option)

        add_opt(wrapper, opt)
        if legacy:
            add_legacy_opt(wrapper, opt)

        return wrapper

    return decorate

It almost works. The thing is that it gives me double legacy messages:

You chose 1!
2 is deprecated!
2 is deprecated!
You chose 2!
3 is not allowed!

If anyone has any ideas, let me know! Thanks :D

Update: Decorating in a different order works:

@option("2", legacy=True)
@option("1")
def do_stuff(opt):
    print(f"You chose {opt}!")
python-decorators metaprogramming python
1 Answers

You've created a wrapper around a wrapper. The wraps decorator actually updates each wrapper with the function attributes from "deeper down," which is actually critical here, but since each wrapper keeps it's own .opts and .legacy (pointing to the same list objects) so each layer of wrappers will print(f"{option} is deprecated!") then return func(option) which is just calling the nested wrapper, which goes through the same checks.

So observe:

In [1]: from functools import wraps
   ...:
   ...:
   ...: def option(opt, legacy=False):
   ...:     def add_opt(f, option):
   ...:         if not hasattr(f, "opts"):
   ...:             f.opts = []
   ...:         f.opts.append(option)
   ...:
   ...:     def add_legacy_opt(f, option):
   ...:         if not hasattr(f, "legacy"):
   ...:             f.legacy = []
   ...:         f.legacy.append(option)
   ...:
   ...:     def decorate(func):
   ...:         @wraps(func)
   ...:         def wrapper(option):
   ...:             if hasattr(wrapper, "opts") and option not in wrapper.opts:
   ...:                 print(f"{option} is not allowed!")
   ...:                 return
   ...:             if hasattr(wrapper, "legacy") and option in wrapper.legacy:
   ...:                 print(f"{option} is deprecated!")
   ...:             return func(option)
   ...:
   ...:         add_opt(wrapper, opt)
   ...:         if legacy:
   ...:             add_legacy_opt(wrapper, opt)
   ...:
   ...:         return wrapper
   ...:
   ...:     return decorate
   ...:

In [2]: @option("2", legacy=True)
   ...: @option("1")
   ...: def do_stuff(opt):
   ...:     print(f"You chose {opt}!")
   ...:

In [3]: vars(do_stuff)
Out[3]:
{'__wrapped__': <function __main__.do_stuff(opt)>,
 'opts': ['1', '2'],
 'legacy': ['2']}

In [4]: vars(do_stuff.__wrapped__)
Out[4]: {'__wrapped__': <function __main__.do_stuff(opt)>, 'opts': ['1', '2']}

In the case where hasattr(wrapper, "opts") and option not in wrapper.opts is true, it just return's without calling func. (as a design aside, this should almost certainly just raise a value error instead).

Decorating in the opposite order "works" because you add legacy last, so only the outermost wrapper has the .legacy attribute.

One solution would be to somehow identify functions that have already been decorated, they are already wrapped, and instead of creating another wrapper, just update the necessary attributes on the wrapper and return the original wrapper.

First, let's move add_opt and add_legacy outside to remove the clutter from within the decorator, there's no good reason for them to be defined inside:

from functools import wraps

def _add_opt(f, option):
    if not hasattr(f, "opts"):
        f.opts = []
    f.opts.append(option)

def _add_legacy_opt(f, option):
    if not hasattr(f, "legacy"):
        f.legacy = []
    f.legacy.append(option)

Then we can do something like the following:

def option(opt, legacy=False):

    def decorate(func):
        if hasattr(func, 'opt') or hasattr(func, 'legacy'): # has been decorated
            wrapper = func # keep old wrapper
        else: # wrap for the first time
            @wraps(func)
            def wrapper(option):
                if hasattr(wrapper, "opts") and option not in wrapper.opts:
                    print(f"{option} is not allowed!")
                    return
                if hasattr(wrapper, "legacy") and option in wrapper.legacy:
                    print(f"{option} is deprecated!")
                return func(option)

        _add_opt(wrapper, opt)
        if legacy:
            _add_legacy_opt(wrapper, opt)

        return wrapper

    return decorate

Honestly, this approach of relying on function attributes seems very brittle. But perhaps you can come up with a more robust approach, but this should at least clarify the problem and give you a solution.

user_50144550
Related