Search code examples
pythonclasssubclassclass-method

Execute classmethod on class initilization to fill class attribute


I am trying to create a Python class with a class attribute the receive a dictionary with all subclasses object.

Below I make a example code within what I am trying to accomplish, but I can't make the classmethod run on class initialization. I have the feeling that I am approaching the problem in the wrong way, so any suggestions will be very welcome.

class Client:

    enable = True
    subclass_dict = {}
        
    def __init__(self, category, limit):
        self.category = category
        self.limit = limit
        if Client.enable == True:
            Client.enable = False
            Client.get_subclasses()
    
    
    @classmethod
    def get_subclasses(cls):
        subclasses = cls.__subclasses__()

        for subclass in subclasses:
            cls.subclass_dict[subclass().category] = subclass()
    
    
    def __str__(self):
        return f'{self.category}'

    def __repr__(self):
        return f'Client: {self.category}'

class SpecialClient(Client):

    category = 'Special'
    limit = '10_000'

    def __init__(self):
        super().__init__(SpecialClient.category, SpecialClient.limit)

class NormalClient(Client):

    category = 'Normal'
    limit = 1_000

    def __init__(self):
        super().__init__(NormalClient.category, NormalClient.limit)

Solution

  • There an easy way to create a subclass "registry" which is effectively what you are doing, but using the generic object.__init_subclass__() that was added in Python 3.6 which makes the finding of subclasses simpler by using it to automatically create a "registry" of them instead of potentially having to check every subclass recursively as a get_subclasses() method would need to do.

    I got the idea of using __init_subclass__() to do this from the Subclass registration section in the PEP 487 -- Simpler customisation of class creation proposal. Since the method will be inherited by all the base class' subclasses, registration will automatically be done for sub-subclasses, too (as opposed to only to direct subclasses) — it completely eliminates the need for a method like get_subclasses().

    I've converted your code to do this and eliminated the get_subclasses() method yours had. There may still be some extraneous stuff in it because I wasn't sure what purpose you might have had for them.

    class Client:
        enable = True
        subclass_dict = {}
    
        def __init__(self, category, limit):
            self.category = category
            self.limit = limit
            if Client.enable == True:
                Client.enable = False
    
        @classmethod
        def __init_subclass__(cls, /, **kwargs):
            super().__init_subclass__(**kwargs)
            cls.subclass_dict[cls.category] = cls  # Add (sub)class to registry.
    
        def __str__(self):
            return f'{self.category}'
    
        def __repr__(self):
            return f'Client: {self.category}'
    
    
    class SpecialClient(Client):
        category = 'Special'
        limit = '10_000'
    
        def __init__(self):
            super().__init__(SpecialClient.category, SpecialClient.limit)
    
    
    class NormalClient(Client):
        category = 'Normal'
        limit = 1_000
    
        def __init__(self):
            super().__init__(NormalClient.category, NormalClient.limit)
    
    
    if __name__ == '__main__':
        from pprint import pprint
        print('Client.subclass_dict:')
        pprint(Client.subclass_dict)
    
        # Verify __str__() and __repr__() were inherited.
        special = SpecialClient()
        print(f'{str(special)=!r}')
        print(f'{repr(special)=!r}')
    
    

    Results printed:

    Client.subclass_dict:
    {'Normal': <class '__main__.NormalClient'>,
     'Special': <class '__main__.SpecialClient'>}
    str(special)='Special'
    repr(special)='Client: Special'