CODEWITHSUNDEEP
pythonpython-3.xdecoratorpython-decorators
Ben
Ben

Reputation: 480

Can I edit a class variable from a method decorator in python?

I have a class as follows:

class Hamburger():
    topping_methods = {}

    @prepare_topping("burger", "onion_ring")
    def cook(self, stuff, temp):
        print("cooking")

h = Hamburger()

I want the prepare_topping decorator to add entries in Hamburger.topping_methods so that it looks like:

{"burger":< cook method reference >, "onion_ring": < cook method reference >}

The kicker is that I need this to happen before the initializer of the class is run (The real use-case is that the dict is used for registering event callbacks in the initializer.) but because I want to access the class variable, it needs to be after the class is defined.

This is as far as I've got with the decorator logic:

def prepare_topping(*args):
    def deco(func):
        # This is when I want to add to class dict, but cannot access method arguments yet
        print(eval(func.__qualname__.split(".")[0]).topping_methods)  # Gets class that method belongs to in a hacky way, but the class is not yet defined
        def wrapper(self, *args):
            print(self.topping_methods)  # Only run if the method is called, which is after __init__
            return func(self, *args)
        return wrapper
    return deco

I realise I will never be able to access the method self argument to achieve this as I want to do stuff regardless of whether or not the method is actually called. Is there a way I can have the decorator run only after the class has been defined? Is there some other way I could achieve this while still using decorators?

Upvotes: 0

Views: 242

Answers (1)

AKX
AKX

Reputation: 169184

You can use a superclass and the __init_subclass__ hook to wire things up:

class CookeryClass:
    topping_methods: dict

    def __init_subclass__(cls, **kwargs):
        cls.topping_methods = {}
        for obj in vars(cls).values():
            if hasattr(obj, "topping_keys"):
                for key in obj.topping_keys:
                    cls.topping_methods[key] = obj


def prepare_topping(*keys):
    def decorator(func):
        func.topping_keys = keys
        return func

    return decorator


class Hamburger(CookeryClass):
    @prepare_topping("burger", "onion_ring")
    def cook(self, stuff, temp):
        print("cooking")

    @prepare_topping("mayonnaise", "pineapples")
    def not_on_a_pizza_surely(self, stuff, temp):
        print("cooking")


print(Hamburger.topping_methods)

This prints out

{
  'burger': <function Hamburger.cook at 0x000001EA49D293A0>,
  'onion_ring': <function Hamburger.cook at 0x000001EA49D293A0>,
  'mayonnaise': <function Hamburger.not_on_a_pizza_surely at 0x000001EA49D29430>,
  'pineapples': <function Hamburger.not_on_a_pizza_surely at 0x000001EA49D29430>,
}

Upvotes: 2

Related Questions