Add option to model meta class to prevent proxy models being polymorphic

SimonFonceca · pgammans · commit 6990167875c7 · 2023-06-08T12:21:11.000+01:00
Fixes jazzband#376 jazzband#390
diff --git a/polymorphic/base.py b/polymorphic/base.py
@@ -57,6 +57,12 @@ class PolymorphicModelBase(ModelBase):
     def __new__(self, model_name, bases, attrs, **kwargs):
         # print; print '###', model_name, '- bases:', bases
 
+        polymorphic__proxy = None
+        if "Meta" in attrs:
+            if hasattr(attrs["Meta"], "polymorphic__proxy"):
+                polymorphic__proxy = attrs["Meta"].polymorphic__proxy
+                del attrs["Meta"].polymorphic__proxy
+
         # Workaround compatibility issue with six.with_metaclass() and custom Django model metaclasses:
         if not attrs and model_name == "NewBase":
             return super().__new__(self, model_name, bases, attrs, **kwargs)
@@ -83,6 +89,11 @@ def __new__(self, model_name, bases, attrs, **kwargs):
         # for __init__ function of this class (monkeypatching inheritance accessors)
         new_class.polymorphic_super_sub_accessors_replaced = False
 
+        if polymorphic__proxy is not None:
+            new_class._meta.polymorphic__proxy = polymorphic__proxy
+        else:
+            new_class._meta.polymorphic__proxy = not new_class._meta.proxy
+
         # determine the name of the primary key field and store it into the class variable
         # polymorphic_primary_key_name (it is needed by query.py)
         for f in new_class._meta.fields:
diff --git a/polymorphic/managers.py b/polymorphic/managers.py
@@ -28,7 +28,7 @@ def from_queryset(cls, queryset_class, class_name=None):
 
     def get_queryset(self):
         qs = self.queryset_class(self.model, using=self._db, hints=self._hints)
-        if self.model._meta.proxy:
+        if not self.model._meta.polymorphic__proxy:
             qs = qs.instance_of(self.model)
         return qs
 
diff --git a/polymorphic/tests/migrations/0001_initial.py b/polymorphic/tests/migrations/0001_initial.py
@@ -1510,6 +1510,84 @@ class Migration(migrations.Migration):
             options={"proxy": True},
             bases=("tests.proxybase",),
         ),
+        migrations.CreateModel(
+            name="AliasProxyChild",
+            fields=[],
+            options={"proxy": True},
+            bases=("tests.proxybase",),
+        ),
+        migrations.CreateModel(
+            name="NonProxyChildAliasProxy",
+            fields=[],
+            options={"proxy": True},
+            bases=("tests.nonproxychild",),
+        ),
+        migrations.CreateModel(
+            name='AliasOfNonProxyChild',
+            fields=[
+            ],
+            options={
+                'proxy': True,
+                'indexes': [],
+                'constraints': [],
+            },
+            bases=('tests.nonproxychild',),
+        ),
+        migrations.CreateModel(
+            name='NonAliasNonProxyChild',
+            fields=[
+            ],
+            options={
+                'proxy': True,
+                'indexes': [],
+                'constraints': [],
+            },
+            bases=('tests.nonproxychild',),
+        ),
+        migrations.CreateModel(
+            name='TradProxyChild',
+            fields=[
+            ],
+            options={
+                'proxy': True,
+                'indexes': [],
+                'constraints': [],
+            },
+            bases=('tests.proxybase',),
+        ),
+        migrations.CreateModel(
+            name='TradProxyOnProxyChild',
+            fields=[
+            ],
+            options={
+                'proxy': True,
+                'indexes': [],
+                'constraints': [],
+            },
+            bases=('tests.proxychild',),
+        ),
+        migrations.CreateModel(
+            name='PolyTradProxyChild',
+            fields=[
+            ],
+            options={
+                'proxy': True,
+                'indexes': [],
+                'constraints': [],
+            },
+            bases=('tests.tradproxychild',),
+        ),
+        migrations.CreateModel(
+            name='ProxyChildAliasProxy',
+            fields=[
+            ],
+            options={
+                'proxy': True,
+                'indexes': [],
+                'constraints': [],
+            },
+            bases=('tests.tradproxychild',),
+        ),
         migrations.CreateModel(
             name="ProxyModelBase",
             fields=[],
diff --git a/polymorphic/tests/models.py b/polymorphic/tests/models.py
@@ -332,7 +332,7 @@ class UUIDPlainC(UUIDPlainB):
     field3 = models.CharField(max_length=10)
 
 
-# base -> proxy
+# base(poly) -> proxy
 
 
 class ProxyBase(PolymorphicModel):
@@ -348,7 +348,59 @@ class NonProxyChild(ProxyBase):
     name = models.CharField(max_length=10)
 
 
-# base -> proxy -> real models
+# A traditional django proxy models. ie proxy'ed class is alias class
+# but in django_polymorphic this is not so.
+#
+# We have model types :-
+# base(poly) / child(poly)   : A concrete polymorphic model 1+ fields
+# base(non ploy)             : A concrete django model 1+ fields
+# proxy(poly)                : A proxy model where it considered different
+#                            : from it superclasses
+# proxy(Traditional Django)  : A proxy model where it is an alias for the
+#                            : underline model
+
+
+# base(poly) -> proxy(poly) -> proxy(Traditional Django)
+class TradProxyOnProxyChild(ProxyChild):
+    class Meta:
+        proxy = True
+        polymorphic__proxy = True
+
+
+# base(poly) -> proxy(Traditional Django)
+class TradProxyChild(ProxyBase):
+    class Meta:
+        proxy = True
+        polymorphic__proxy = True
+
+# base(poly) -> proxy(Traditional Django) -> proxy(poly)
+# Not really helpful model as reduces to base(poly) -> proxy(poly)
+
+# base(poly) -> child(poly) -> proxy(Traditional Django)
+class AliasOfNonProxyChild(NonProxyChild):
+    class Meta:
+        proxy = True
+        polymorphic__proxy = True
+
+
+# base(poly) -> proxy(Traditional Django) -> proxy(poly)
+class ProxyChildAliasProxy(TradProxyChild):
+    class Meta:
+        proxy = True
+
+
+# base(poly) -> proxy(poly)
+class AliasProxyChild(ProxyBase):
+    class Meta:
+        proxy = True
+        polymorphic__proxy = True
+
+
+# base(poly) -> proxy(poly)
+class NonAliasNonProxyChild(NonProxyChild):
+    class Meta:
+        proxy = True
+        polymorphic__proxy = False
 
 
 class ProxiedBase(ShowFieldTypeAndContent, PolymorphicModel):
diff --git a/polymorphic/tests/test_orm.py b/polymorphic/tests/test_orm.py
@@ -11,6 +11,7 @@
 from polymorphic.managers import PolymorphicManager
 from polymorphic.models import PolymorphicTypeInvalid, PolymorphicTypeUndefined
 from polymorphic.tests.models import (
+    AliasProxyChild,
     ArtProject,
     Base,
     BlogA,
@@ -87,6 +88,13 @@
     UUIDPlainC,
     UUIDProject,
     UUIDResearchProject,
+
+    NonAliasNonProxyChild,
+    TradProxyOnProxyChild,
+    TradProxyChild,
+    AliasOfNonProxyChild,
+    ProxyChildAliasProxy,
+
 )
 
 
@@ -875,6 +883,61 @@ def test_queryset_on_proxy_model_does_not_return_superclasses(self):
         self.assertEqual(5, ProxyBase.objects.count())
         self.assertEqual(3, ProxyChild.objects.count())
 
+    def test_queryset_on_polymorphic_proxy_model_returns_superclasses(self):
+        ProxyBase.objects.create(some_data="Base1")
+        AliasProxyChild.objects.create(some_data="ProxyChild1")
+        AliasProxyChild.objects.create(some_data="ProxyChild2")
+        ProxyChild.objects.create(some_data="PolyChild1")
+        NonAliasNonProxyChild.objects.create(some_data="SubChild1")
+        NonAliasNonProxyChild.objects.create(some_data="SubChild2")
+        NonProxyChild.objects.create(some_data="NonProxChild1", name="t1")
+
+        with self.subTest(" superclasses"):
+            self.assertEqual(7, ProxyBase.objects.count())
+            self.assertEqual(7, AliasProxyChild.objects.count())
+        with self.subTest("only compete classes"):
+            # Non proxy models should not return the proxy siblings
+            self.assertEqual(1, ProxyChild.objects.count())
+            self.assertEqual(2, NonAliasNonProxyChild.objects.count())
+            self.assertEqual(3, NonProxyChild.objects.count())
+
+    def test_polymorphic_proxy_object_has_different_ctype_from_base(self):
+        obj1 = ProxyBase.objects.create(some_data="Base1")
+        obj2 = AliasProxyChild.objects.create(some_data="ProxyChild1")
+        obj1_ctype = ContentType.objects.get_for_model(
+            obj1, for_concrete_model=False)
+        obj2_ctype = ContentType.objects.get_for_model(
+            obj2, for_concrete_model=False)
+        self.assertNotEqual(obj1_ctype, obj2_ctype)
+
+    def test_can_create_django_style_proxy_classes_alias(self):
+        ProxyBase.objects.create(some_data="Base1")
+        TradProxyChild.objects.create(some_data="Base2")
+        self.assertEqual(2, ProxyBase.objects.count())
+        self.assertEqual(2, TradProxyChild.objects.count())
+        TradProxyOnProxyChild.objects.create()
+
+    def test_convert_back_to_django_style_from_polymorphic(self):
+        ProxyBase.objects.create(some_data="Base1")
+        ProxyChild.objects.create(some_data="Base1")
+        TradProxyOnProxyChild.objects.create(some_data="Base3")
+        
+        self.assertEqual(3, ProxyBase.objects.count())
+        self.assertEqual(2, ProxyChild.objects.count())
+        self.assertEqual(3, TradProxyOnProxyChild.objects.count())
+
+    def test_convert_back_to_django_style_from_polymorphic_stops_at_concrete(self):
+        ProxyBase.objects.create(some_data="Base1")
+        NonProxyChild.objects.create(some_data="Base1")
+        AliasOfNonProxyChild.objects.create(some_data="Base1")
+
+        self.assertEqual(3, ProxyBase.objects.count())
+        self.assertEqual(2, NonProxyChild.objects.count())
+        self.assertEqual(2, AliasOfNonProxyChild.objects.count())
+
+    def test_revert_back_to_polymorphic_proxy(self):
+        self.assertFalse(ProxyChildAliasProxy._meta.polymorphic__proxy)
+
     def test_proxy_get_real_instance_class(self):
         """
         The call to ``get_real_instance()`` also checks whether the returned model is of the correct type.
@@ -884,12 +947,12 @@ def test_proxy_get_real_instance_class(self):
         name = "Item1"
         nonproxychild = NonProxyChild.objects.create(name=name)
 
-        pb = ProxyBase.objects.get(id=1)
+        pb = ProxyBase.objects.get(id=nonproxychild.pk)
         self.assertEqual(pb.get_real_instance_class(), NonProxyChild)
         self.assertEqual(pb.get_real_instance(), nonproxychild)
         self.assertEqual(pb.name, name)
 
-        pbm = NonProxyChild.objects.get(id=1)
+        pbm = NonProxyChild.objects.get(id=nonproxychild.pk)
         self.assertEqual(pbm.get_real_instance_class(), NonProxyChild)
         self.assertEqual(pbm.get_real_instance(), nonproxychild)
         self.assertEqual(pbm.name, name)
