Added get_base_polymorphic_model() to detect the common base class for a polymorphic model.

polymorphic/tests/test_utils.py

@@ -1,8 +1,8 @@
 from django.test import TransactionTestCase
 
-from polymorphic.models import PolymorphicTypeUndefined
-from polymorphic.tests.models import Model2A, Model2B, Model2C, Model2D
-from polymorphic.utils import reset_polymorphic_ctype, sort_by_subclass
+from polymorphic.models import PolymorphicTypeUndefined, PolymorphicModel
+from polymorphic.tests.models import Model2A, Model2B, Model2C, Model2D, Enhance_Inherit, Enhance_Base
+from polymorphic.utils import reset_polymorphic_ctype, sort_by_subclass, get_base_polymorphic_model
 
 
 class UtilsTests(TransactionTestCase):
@@ -38,3 +38,39 @@ class UtilsTests(TransactionTestCase):
             ],
             transform=lambda o: o.__class__,
         )
+
+    def test_get_base_polymorphic_model(self):
+        """
+        Test that finding the base polymorphic model works.
+        """
+        # Finds the base from every level (including lowest)
+        self.assertIs(get_base_polymorphic_model(Model2D), Model2A)
+        self.assertIs(get_base_polymorphic_model(Model2C), Model2A)
+        self.assertIs(get_base_polymorphic_model(Model2B), Model2A)
+        self.assertIs(get_base_polymorphic_model(Model2A), Model2A)
+
+        # Properly handles multiple inheritance
+        self.assertIs(get_base_polymorphic_model(Enhance_Inherit), Enhance_Base)
+
+        # Ignores PolymorphicModel itself.
+        self.assertIs(get_base_polymorphic_model(PolymorphicModel), None)
+
+    def test_get_base_polymorphic_model_skip_abstract(self):
+        """
+        Skipping abstract models that can't be used for querying.
+        """
+        class A(PolymorphicModel):
+            class Meta:
+                abstract = True
+
+        class B(A):
+            pass
+
+        class C(B):
+            pass
+
+        self.assertIs(get_base_polymorphic_model(A), None)
+        self.assertIs(get_base_polymorphic_model(B), B)
+        self.assertIs(get_base_polymorphic_model(C), B)
+
+        self.assertIs(get_base_polymorphic_model(C, allow_abstract=True), A)