Add inheritance support

lokhman · Dec 18, 2020 · d0dcbd8 · d0dcbd8
1 parent e3db7b3
commit d0dcbd8
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 4 deletions.
diff --git a/django_better_choices/__init__.py b/django_better_choices/__init__.py
@@ -132,13 +132,19 @@ def __new__(cls, __name: Optional[str] = None, **kwargs: Any) -> Union[Type[Choi
             return tuple(v.__choice_entry__ for _, v in cls.__iter_items(**kwargs))
         return cast(Type[Choices], type(cls.__name__ if __name is None else __name, (Choices,), kwargs))
 
-    def __init_subclass__(cls, **kwargs: Any):
+    def __init_subclass__(cls, _subset: bool = False, **kwargs: Any):
         super().__init_subclass__(**kwargs)
 
         cls.__keys = {}
         cls.__values = {}
 
-        for key, value in cls.__dict__.items():
+        __dict__ = {}
+        if not _subset:
+            for base in reversed(cls.__mro__[1:-2]):
+                __dict__.update(base.__dict__)
+        __dict__.update(cls.__dict__)
+
+        for key, value in __dict__.items():
             if key.startswith("_"):
                 continue
 
@@ -256,15 +262,25 @@ def extract(cls, __key: str, *keys: str, name: str = "Subset") -> Type[Choices]:
         """Extract values from choices class and return a new subset."""
         return cast(
             Type[Choices],
-            type(f"{cls.__name__}.{name}", (cls,), {k: cls.__values[k] for k in (__key, *keys)}),
+            type(
+                f"{cls.__name__}.{name}",
+                (cls,),
+                {k: cls.__values[k] for k in (__key, *keys)},
+                _subset=True,
+            ),
         )
 
     @classmethod
     def exclude(cls, __key: str, *keys: str, name: str = "Subset") -> Type[Choices]:
         """Exclude values from choices class and return remaining values as a new subset."""
         return cast(
             Type[Choices],
-            type(f"{cls.__name__}.{name}", (cls,), {k: v for v, k in cls.__keys.items() if k not in {__key, *keys}}),
+            type(
+                f"{cls.__name__}.{name}",
+                (cls,),
+                {k: v for v, k in cls.__keys.items() if k not in {__key, *keys}},
+                _subset=True,
+            ),
         )
 
     Value = _ChoicesValue

diff --git a/tests.py b/tests.py
@@ -244,6 +244,24 @@ def test_custom_methods(self):
         self.assertEqual(("DISPLAY 3", "DISPLAY 5"), TestChoices.SUBSET2.get_upper_displays())
         self.assertEqual(("DISPLAY 1", "DISPLAY 7"), TestChoices.extract("VAL1", "VAL7").get_upper_displays())
 
+    def test_inheritance(self):
+        class TestNextChoices(TestChoices):
+            VAL3 = Choices.Value("Display 3", value="val3")
+            VAL8 = "Display 8"
+
+        class TestFinalChoices(TestNextChoices):
+            VAL9 = "Display 9"
+            SUBSET3 = Choices.Subset("VAL2", "VAL8", "VAL9")
+
+        self.assertEqual("Display 1", TestFinalChoices.VAL1.display)
+        self.assertEqual("val3", TestFinalChoices.VAL3)
+        self.assertEqual("val9", TestFinalChoices.VAL9)
+        self.assertEqual(
+            ("val1", "val2", "val3", "val4", "val5", (1, 2, 3), 7, "val8", "val9"),
+            TestFinalChoices.values(),
+        )
+        self.assertEqual(("DISPLAY 2", "DISPLAY 8", "DISPLAY 9"), TestFinalChoices.SUBSET3.get_upper_displays())
+
     def test_copy(self):
         self.assertEqual("val1", copy.copy(TestChoices.VAL1))
         self.assertEqual("value-3", copy.copy(TestChoices.VAL3))