Add SubClasses registration to BaseEntity (#532)

LeMyst · Apr 23, 2023 · f265520 · f265520
1 parent 63dac0c
commit f265520
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/wikibaseintegrator/entities/baseentity.py b/wikibaseintegrator/entities/baseentity.py
@@ -2,7 +2,7 @@
 
 import logging
 from copy import copy
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
 from wikibaseintegrator import wbi_fastrun
 from wikibaseintegrator.datatypes import BaseDataType
@@ -20,6 +20,7 @@
 
 class BaseEntity:
     ETYPE = 'base-entity'
+    subclasses: List[Type[BaseDataType]] = []
 
     def __init__(self, api: Optional['WikibaseIntegrator'] = None, title: Optional[str] = None, pageid: Optional[int] = None, lastrevid: Optional[int] = None,
                  type: Optional[str] = None, id: Optional[str] = None, claims: Optional[Claims] = None, is_bot: Optional[bool] = None, login: Optional[_Login] = None):
@@ -39,6 +40,11 @@ def __init__(self, api: Optional['WikibaseIntegrator'] = None, title: Optional[s
         self.id = id
         self.claims = claims or Claims()
 
+    # Allow registration of subclasses of BaseEntity into BaseEntity.subclasses
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls.subclasses.append(cls)
+
     @property
     def api(self) -> WikibaseIntegrator:
         return self.__api