gath gath - 1 month ago 11
Python Question

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.

Answer

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