feat: add support to convert from normal manager/qs to poly versions

Author: pgammans

docs/advanced.rst

@@ -177,6 +177,13 @@ About Queryset Methods
             relation = models.ForeignKey(BasePolyModel, on_delete=models.CASCADE)
             objects = models.Manager.from_queryset(PolymorphicRelatedQuerySet)()
 
+    or by converting a models queryset using
+
+        class NonPolyModel(models.Model):
+            relation = models.ForeignKey(BasePolyModel, on_delete=models.CASCADE)
+            objects = models.Manager.from_queryset(QuerySet)()
+
+    ``convert_to_polymorphic_queryset(NonPolyModel.objects).filter(...)``
 
     To select related fields the model name comes after the field name and set the
     field.

polymorphic/query.py

@@ -9,7 +9,7 @@
 
 from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import FieldDoesNotExist
-from django.db.models import FilteredRelation
+from django.db.models import FilteredRelation, Manager
 from django.db.models.constants import LOOKUP_SEP
 from django.db.models.query import ModelIterable, Q, QuerySet
 from django.db.models.query import BaseIterable, RelatedPopulator
@@ -1015,3 +1015,36 @@ def _get_real_instances(self, base_result_objects):
 
 class PolymorphicRelatedQuerySet(PolymorphicRelatedQuerySetMixin, QuerySet):
     pass
+
+
+def convert_to_polymorphic_queryset(qs):
+    "Convert a queryset to one that support polymorphic evaluation"
+
+    if isinstance(qs, Manager):
+        qs = qs.get_queryset()
+
+    if issubclass(qs.__class__, PolymorphicQuerySetMixin):
+        return qs
+
+    assert issubclass(QuerySet, qs.__class__), (
+        "PolymorphicModel: cannot guarantee conversion of"
+        f" {qs.__class__} to polymorphic queryset")
+
+    class RelatedPolyQuerySet(PolymorphicRelatedQuerySetMixin, qs.__class__):
+        @classmethod
+        def _convert_to(cls, qs):
+            c = cls(
+                model=qs.model,
+                query=qs.query.chain(),
+                using=qs._db,
+                hints=qs._hints,
+            )
+            c._sticky_filter = qs._sticky_filter
+            c._for_write = qs._for_write
+            c._prefetch_related_lookups = qs._prefetch_related_lookups[:]
+            c._known_related_objects = qs._known_related_objects
+            c._fields = qs._fields
+            return c
+
+    poly_qs = RelatedPolyQuerySet._convert_to(qs)
+    return poly_qs

polymorphic/tests/migrations/0001_initial.py

@@ -2344,4 +2344,22 @@ class Migration(migrations.Migration):
                 ),
             ],
         ),
+        migrations.CreateModel(
+            name="VanillaPlainModel",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True, primary_key=True, serialize=False, verbose_name="ID"
+                    ),
+                ),
+                (
+                    "relation",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE, to="tests.parentmodel"
+                    ),
+                ),
+            ],
+        ),
+
     ]

polymorphic/tests/models.py

@@ -554,10 +554,13 @@ class PlainModel(models.Model):
     relation = models.ForeignKey(ParentModel, on_delete=models.CASCADE)
     objects = models.Manager.from_queryset(PolymorphicRelatedQuerySet)()
 
+class VanillaPlainModel(models.Model):
+    relation = models.ForeignKey(ParentModel, on_delete=models.CASCADE)
 
 class RefPlainModel(models.Model):
     plainobj = models.ForeignKey(PlainModel, on_delete=models.CASCADE)
-    objects = models.Manager.from_queryset(PolymorphicRelatedQuerySet)()
+    objects = models.Manager.from_queryset(QuerySet)()
+    poly_objects = models.Manager.from_queryset(PolymorphicRelatedQuerySet)()
 
 
 class PlainModelWithM2M(models.Model):

polymorphic/tests/test_orm.py

@@ -10,6 +10,7 @@
 from django.db.utils import IntegrityError
 from django.test import TransactionTestCase
 from polymorphic import query_translate
+from polymorphic.query import convert_to_polymorphic_queryset
 from polymorphic.managers import PolymorphicManager
 from polymorphic.models import PolymorphicTypeInvalid, PolymorphicTypeUndefined
 from polymorphic.tests.models import (
@@ -101,6 +102,7 @@
     UUIDPlainC,
     UUIDProject,
     UUIDResearchProject,
+    VanillaPlainModel,
 )
 
 
@@ -1428,6 +1430,74 @@ def test_select_related_on_poly_classes_simple(self):
         self.assertIsInstance(obj_list[2].relation, AltChildModel)
         self.assertIsInstance(obj_list[3].relation, AltChildModel)
 
+    def test_we_can_upgrade_a_query_set_to_polymorphic_supports_already_ploy_qs(self):
+        base_qs = RefPlainModel.poly_objects.get_queryset()
+        self.assertIs(convert_to_polymorphic_queryset(base_qs), base_qs)
+
+    def test_we_can_upgrade_a_query_set_to_polymorphic_supports_non_ploy_qs_on_ploy_object(self):
+        base_qs = RefPlainModel.objects.get_queryset()
+        self.assertIsNot(convert_to_polymorphic_queryset(base_qs), base_qs)
+        self.assertIsInstance(convert_to_polymorphic_queryset(base_qs), PolymorphicRelatedQuerySetMixin)
+
+    def test_we_can_upgrade_a_query_set_to_polymorphic_supports_non_ploy_managers_on_ploy_object(self):
+        base_qs = RefPlainModel.objects
+        self.assertIsNot(convert_to_polymorphic_queryset(base_qs), base_qs)
+        self.assertIsInstance(convert_to_polymorphic_queryset(base_qs), PolymorphicRelatedQuerySetMixin)
+
+    def test_we_can_upgrade_a_query_set_to_polymorphic(self):
+        # can we fetch the related object but only the minimal 'common' values
+        plain_a_obj_1 = PlainA.objects.create(field1="f1")
+        plain_a_obj_2 = PlainA.objects.create(field1="f2")
+        extra_obj = ModelExtraExternal.objects.create(topic="t1")
+        obj_p = ParentModel.objects.create(name="p1")
+        obj_c = ChildModel.objects.create(name="c1", other_name="c1name", link_on_child=extra_obj)
+        obj_ac1 = AltChildModel.objects.create(
+            name="ac1", other_name="ac1name", link_on_altchild=plain_a_obj_1
+        )
+        obj_ac2 = AltChildModel.objects.create(
+            name="ac2", other_name="ac2name", link_on_altchild=plain_a_obj_2
+        )
+        obj_p_1 = VanillaPlainModel.objects.create(relation=obj_p)
+        obj_p_2 = VanillaPlainModel.objects.create(relation=obj_c)
+        obj_p_3 = VanillaPlainModel.objects.create(relation=obj_ac1)
+        obj_p_4 = VanillaPlainModel.objects.create(relation=obj_ac2)
+
+        with self.assertNumQueries(1):
+            # pos 3 if i cannot do optimized select_related
+            obj_list = list(
+                VanillaPlainModel.objects.order_by("pk")
+            )
+
+        with self.assertNumQueries(7):
+            self.assertEqual(obj_list[0].relation.name, "p1")
+            self.assertEqual(obj_list[1].relation.name, "c1")
+            self.assertEqual(obj_list[2].relation.name, "ac1")
+            self.assertEqual(obj_list[3].relation.name, "ac2")
+
+        with self.assertNumQueries(1):
+            # pos 3 if i cannot do optimized select_related
+            obj_list = list(
+                convert_to_polymorphic_queryset(
+                    VanillaPlainModel.objects
+                ).select_related(
+                    "relation",
+                    "relation__childmodel",
+                    "relation__altchildmodel",
+                )
+                .order_by("pk")
+            )
+
+        with self.assertNumQueries(0):
+            self.assertEqual(obj_list[0].relation.name, "p1")
+            self.assertEqual(obj_list[1].relation.name, "c1")
+            self.assertEqual(obj_list[2].relation.name, "ac1")
+            self.assertEqual(obj_list[3].relation.name, "ac2")
+
+        self.assertIsInstance(obj_list[0].relation, ParentModel)
+        self.assertIsInstance(obj_list[1].relation, ChildModel)
+        self.assertIsInstance(obj_list[2].relation, AltChildModel)
+        self.assertIsInstance(obj_list[3].relation, AltChildModel)
+
     def test_select_related_on_poly_classes_indirect_related(self):
         # can we fetch the related object but only the minimal 'common' values
         plain_a_obj_1 = PlainA.objects.create(field1="f1")
@@ -1457,7 +1527,7 @@ def test_select_related_on_poly_classes_indirect_related(self):
         with self.assertNumQueries(1):
             # pos 3 if i cannot do optimized select_related
             obj_list = list(
-                RefPlainModel.objects.select_related(
+                RefPlainModel.poly_objects.select_related(
                     # "plainobj__relation",
                     "plainobj__relation",
                     "plainobj__relation__childmodel__link_on_child",
@@ -1504,7 +1574,7 @@ def test_select_related_on_poly_classes_indirect_related(self):
         with self.assertNumQueries(1):
             # pos 3 if i cannot do optimized select_related
             obj_list = list(
-                RefPlainModel.objects.select_related(
+                RefPlainModel.poly_objects.select_related(
                     # "plainobj__relation",
                     "plainobj__relation",
                     "plainobj__relation__childmodel__link_on_child",
@@ -1554,7 +1624,7 @@ def test_select_related_fecth_all_poly_classes_indirect_related(self):
         with self.assertNumQueries(1):
             # pos 3 if i cannot do optimized select_related
             obj_list = list(
-                RefPlainModel.objects.select_related(
+                RefPlainModel.poly_objects.select_related(
                     # "plainobj__relation",
                     "plainobj__relation",
                     "plainobj__relation__*",