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)
CodePudding user response:
There an easy way to create a subclass "registry" which is effectively what you are doig, 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)
Results:
Client.subclass_dict:
{'Normal': <class '__main__.NormalClient'>,
'Special': <class '__main__.SpecialClient'>}