CODEWITHSUNDEEP
pythonpython-decorators
gath
gath

Reputation: 25492

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.

Upvotes: 2

Views: 541

Answers (4)

Serge Ballesta
Serge Ballesta

Reputation: 149125

If you want to register the instances of annotated subclasses, the annotation must return something that will create the instances and register them.

By the way, if you use a simple dictionary, you will only register the last created instance of each subclass. If you want to register all instances, you should better use a defaultdict(list)

Here is an example of code:

class Model(object):
    models = collections.defaultdict(list)


def register(cls):
    def registrar(*args, **kwargs):
        # first create the instance with the passed parameters
        instance = cls(*args, **kwargs) 
        # then register the instance and return it
        Model.models[cls.__name__].append(instance)
        return instance
    return registrar

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

This code can register different instances:

>>> p1 = PaperModel('A4')
>>> p1
<__main__.PaperModel object at 0x02AA0FB0>
>>> p2 = PaperModel('A3')
>>> p2
<__main__.PaperModel object at 0x02AA0F10>
>>> Model.models
defaultdict(<type 'list'>, {'PaperModel': [<__main__.PaperModel object at 0x02AA0FB0>, <__main__.PaperModel object at 0x02AA0F10>]})

Alternatively, you could not use an annotation at all, and use the special __new__ method of the base class to obtain same result:

class Model(object):
    models = collections.defaultdict(list)
    def __new__(cls, *args, **kwargs):
        inst = super(Model, cls).__new__(cls)
        inst.__init__(*args, **kwargs)
        Model.models[cls.__name__].append(inst)
        return inst


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

This code is able to register instance too:

>>> p1 = PaperModel('A4')
>>> p2 = PaperModel('A3')
>>> Model.models
defaultdict(<type 'list'>, {'PaperModel': [<__main__.PaperModel object at 0x02AA02D0>, <__main__.PaperModel object at 0x02AA03F0>]})
>>> Model.models['PaperModel'][0] is p1
True
>>> Model.models['PaperModel'][1] is p2
True

Upvotes: 4

Moinuddin Quadri
Moinuddin Quadri

Reputation: 48100

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

Upvotes: 2

Moses Koledoye
Moses Koledoye

Reputation: 78564

You should register the instance only when it's instantiated not before:

def register(cls):
    def wrapper(*args, **kwargs):
        inst = cls(*args, **kwargs)
        Model.models[cls.__name__] = inst
        return inst
    return wrapper

Note that you also need to return the class from the decorator else your decorator'll return None you'll get a TypeError when trying to use the decorator on the class.

wm = WoodenModel('log')
print(wm)
# <__main__.WoodenModel object at 0x033A6730>
print(Model.models)
# {'WoodenModel': <__main__.WoodenModel object at 0x033A6730>}

Upvotes: 3

John Zwinck
John Zwinck

Reputation: 249502

You can register the type rather than an instance (which you don't know how to construct, you don't have the right arguments available):

Model.models[cls.__name__] = cls # note no parens

Upvotes: 0

Related Questions