Merge pull request jazzband#739 from bckohan/recursion-bug

Fix infinite recursion bug when using only()

Author: bckohan

codecov.yml

@@ -0,0 +1,2 @@
+fixes:
+  - "src/::"

src/polymorphic/query.py

@@ -219,7 +219,8 @@ def only(self, *fields):
         when we're retrieving the real instance (since we'll need to translate
         them again, as the model will have changed).
         """
-        new_fields = [translate_polymorphic_field_path(self.model, a) for a in fields]
+        new_fields = {translate_polymorphic_field_path(self.model, a) for a in fields}
+        new_fields.add("polymorphic_ctype_id")
         clone = super().only(*new_fields)
         clone._polymorphic_add_immediate_loading(fields)
         return clone

src/polymorphic/tests/migrations/0001_initial.py

@@ -1,4 +1,4 @@
-# Generated by Django 4.2 on 2025-12-14 01:14
+# Generated by Django 4.2 on 2025-12-16 21:55
 
 from django.conf import settings
 from django.db import migrations, models
@@ -13,8 +13,8 @@ class Migration(migrations.Migration):
     initial = True
 
     dependencies = [
-        ('contenttypes', '0002_remove_content_type_name'),
         ('auth', '0012_alter_user_first_name_max_length'),
+        ('contenttypes', '0002_remove_content_type_name'),
     ]
 
     operations = [
@@ -893,6 +893,18 @@ class Migration(migrations.Migration):
             },
             bases=(polymorphic.showfields.ShowFieldType, models.Model),
         ),
+        migrations.CreateModel(
+            name='RecursionBug',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('polymorphic_ctype', models.ForeignKey(editable=False, null=True, on_delete=django.db.models.deletion.CASCADE, related_name='polymorphic_%(app_label)s.%(class)s_set+', to='contenttypes.contenttype')),
+                ('status', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='recursions', to='tests.plaina')),
+            ],
+            options={
+                'abstract': False,
+                'base_manager_name': 'objects',
+            },
+        ),
         migrations.CreateModel(
             name='RankedAthlete',
             fields=[

src/polymorphic/tests/models.py

@@ -770,3 +770,14 @@ class RankedAthlete(models.Model):
     choiceAthlete = models.ForeignKey("ChoiceBlank", on_delete=models.CASCADE)
     bet = models.ForeignKey("BetMultiple", on_delete=models.CASCADE)
     rank = models.IntegerField()
+
+
+class RecursionBug(PolymorphicModel):
+    status = models.ForeignKey(PlainA, on_delete=models.CASCADE, related_name="recursions")
+
+    def __init__(self, *args, **kwargs):
+        """
+        https://github.com/jazzband/django-polymorphic/issues/334
+        """
+        super().__init__(*args, **kwargs)
+        self.old_status_id = self.status_id

src/polymorphic/tests/test_orm.py

@@ -1929,3 +1929,19 @@ def test_through_model_updates(self):
         assert list(bet.answer.order_by("rankedathlete__rank")) == [a2, a1]
         assert RankedAthlete.objects.get(pk=ra.pk).rank == 99
         assert RankedAthlete.objects.get(pk=ra2.pk).rank == 0
+
+    def test_infinite_recursion_with_only(self):
+        """
+        https://github.com/jazzband/django-polymorphic/issues/334
+        """
+        from polymorphic.tests.models import RecursionBug
+
+        draft = PlainA.objects.create(field1="draft")
+        closed = PlainA.objects.create(field1="closed")
+
+        assert isinstance(closed.recursions, PolymorphicManager)
+
+        item = RecursionBug.objects.create(status=draft)
+        RecursionBug.objects.filter(id=item.id).update(status=closed)
+        item.refresh_from_db(fields=("status",))
+        assert item.status == closed