Search code examples
pythonpython-decorators

Python decorating class


I'm trying to decorate a class with arguments but cannot get it to work:

This is the decorator:

 def message(param1, param2):
   def get_message(func):
        func.__init__(param1,param2)

   return get_message

class where I want to put the decorator

@message(param1="testing1", param2="testing2")
class SampleClass(object):
   def __init__(self):
     pass

But this is not working , I am getting an error when running this. Does anyone know what the problem ?, I am trying to create a decorator to initialise classes with some values.


Solution

  • I'm having trouble figuring out what you're trying to do. If you want to decorate a class with a decorator that takes arguments, one way to do it is like this.

    2020-09-03:

    Thank you Maria-Ines Carrera for pointing out that the original code doesn't actually handle classes inheriting from other classes correctly and user2357112 supports Monica for proposing a solution that does work.

    # function returning a decorator, takes arguments
    def message(param1, param2):
        # this does the actual heavy lifting of decorating the class
        # it takes the original class, modifies it in place, and returns
        # the same class
        def wrapper(wrapped):
            the_init = wrapped.__init__
    
            def new_init(self):
                self.param1 = param1
                self.param2 = param2
                the_init(self)
    
            def get_message(self):
                return "message %s %s" % (self.param1, self.param2)
    
            wrapped.__init__ = new_init
            wrapped.get_message = get_message
    
            return wrapped
        return wrapper
    
    
    class Pizza(object):
        def __init__(self):
            print "Pizza initialization"
    
    
    @message("param1", "param2")
    class Pizza2(Pizza):
        def __init__(self):
            print "Pizza2 initialization"
            super(Pizza2, self).__init__()
    
    pizza_with_message = Pizza2()
    
    # prints "message param1 param2"
    print pizza_with_message.get_message()
    

    This prints the following:

    Pizza2 initialization
    Pizza initialization
    message param1 param2