Search code examples
pythonpython-decorators

Decorating a class whose constructor has arguments


How do I decorate a class whose constructor has arguments? This is my code;

# Base Class
class Model(object):
    models = {}
    def __init__(self):
        pass

# decorator
def register(cls):
    Model.models[cls.__name__] = cls()

#Subclasses

@register
class PaperModel(Model):
    def __init__(self, paper):
        self.paper = paper

@register
class WoodenModel(Model):
    def __init__(self, wood):
        self.wood = wood

The idea is to register instances of the subclass in the dict inside the base class.

When I run the code I get the following error

    Model.models[cls.__name__] = cls()
TypeError: __init__() takes exactly 2 arguments (1 given)

However, if I remove the arguments in the subclass (PaperModel & WoodenModel) constructors the code works.


Solution

  • You need instance i.e. object of the class (not the class itself) to be registered in the base class. Cleaner way would be to register the instances in the __init__ of parent class, as the sub classes are derived from it. There is no point in creating the specific decorator for this. Decorators are for generic use, for example if Parent class was also dynamic. Your decorator should be registering things like: <SomeClass>.models[<some_sub_class>]

    Below is the sample code for registering in the parent __init__:

    # update the entry the __init__() of parent class
    class Model(object):
        models = {}
        def __init__(self):
            Model.models[self.__class__.__name__] = self # register instance
    
    class WoodenModel(Model):
        def __init__(self, wood):
            self.wood = wood
            super(self.__class__, self).__init__()  # Make a call to parent's init()
    
    # Create a object
    wooden_model_obj = WoodenModel(123)
    
    print wooden_model_obj
    # prints: <__main__.WoodenModel object at 0x104b3e990>
    #                                            ^
    
    print Model.models
    # prints: {'WoodenModel': <__main__.WoodenModel object at 0x104b3e990>}
    #                                                            ^
    
    # Both referencing same object
    

    In case you want a generic decorator to achieve this, assuming:

    • the child class has only one parent
    • the attribute storing the value is as models in perent class

    The sample decorator will be as:

    def register(cls):
        def register_wrapper(*args, **kwargs):
            obj = cls(*args, **kwargs)
            obj.__class__.__bases__[0].models[cls.__name__] = obj
            # "obj.__class__.__bases__[0]" will return first inherited parent class
            return obj
        return register_wrapper