Refactor variant dto, adding variant_count and removing it from combo…

…s dto

backend/spellbook/management/commands/export_variants.py

@@ -28,6 +28,7 @@ def prepare_variant_alias(variant_alias: VariantAlias):
 class Command(AbstractCommand):
     name = 'export_variants'
     help = 'Exports variants to a JSON file'
+    batch_size = 5000
 
     def add_arguments(self, parser):
         super().add_arguments(parser)
@@ -45,11 +46,14 @@ def add_arguments(self, parser):
 
     def run(self, *args, **options):
         self.log('Fetching variants from db...')
+        variants_source = list[Variant]()
         with transaction.atomic(durable=True):
-            variants_source = list[Variant](VariantSerializer.prefetch_related(Variant.objects.filter(status__in=Variant.public_statuses() + Variant.preview_statuses())))
+            variants_query = VariantSerializer.prefetch_related(Variant.objects.filter(status__in=Variant.public_statuses() + Variant.preview_statuses()))
+            for i in range(0, variants_query.count(), self.batch_size):
+                variants_source.extend(variants_query[i:i + self.batch_size])
         self.log('Fetching variants from db...done', self.style.SUCCESS)
         self.log('Updating variants preserialized representation...')
-        Variant.objects.bulk_serialize(objs=variants_source, serializer=VariantSerializer, batch_size=5000)  # type: ignore
+        Variant.objects.bulk_serialize(objs=variants_source, serializer=VariantSerializer, batch_size=self.batch_size)  # type: ignore
         self.log('Updating variants preserialized representation...done', self.style.SUCCESS)
         self.log('Fetching variant aliases from db...')
         with transaction.atomic(durable=True):

backend/spellbook/management/commands/update_cards.py

@@ -14,16 +14,17 @@ def run(self, *args, **options):
         scryfall_name_db = scryfall()
         self.log('Fetching Scryfall and EDHREC datasets...done')
         self.log('Updating cards...')
-        cards_to_update = list(
-            Card.objects.annotate(
+        cards_to_update = list(Card.objects.all())
+        cards_count: dict[int, int] = {
+            i: c
+            for i, c in Card.objects.annotate(
                 updated_variant_count=Count('used_in_variants', distinct=True, filter=Q(used_in_variants__status__in=Variant.public_statuses())),
-            ),
-        )
-        for card_to_update in cards_to_update:
-            card_to_update.variant_count = card_to_update.updated_variant_count  # type: ignore
+            ).values_list('id', 'updated_variant_count')
+        }
         cards_to_save = update_cards(
             cards_to_update,
             scryfall_name_db,
+            cards_count,
             log=lambda x: self.log(x),
             log_warning=lambda x: self.log(x, self.style.WARNING),
             log_error=lambda x: self.log(x, self.style.ERROR),

backend/spellbook/management/commands/update_variants.py

@@ -1,6 +1,6 @@
 from spellbook.models import Variant
 from spellbook.models.combo import Combo
-from django.db.models import Subquery, OuterRef, Count
+from django.db.models import Q, Subquery, OuterRef, Count
 from django.db.models.functions import Coalesce
 from ..abstract_command import AbstractCommand
 from ..edhrec import edhrec, update_variants
@@ -9,22 +9,34 @@
 class Command(AbstractCommand):
     name = 'update_variants'
     help = 'Updates variants using cards and EDHREC data'
+    batch_size = 5000
 
     def run(self, *args, **options):
         self.log('Fetching EDHREC dataset...')
         edhrec_variant_db = edhrec()
         self.log('Updating variants...')
         variants_query = Variant.objects.prefetch_related('uses', 'cardinvariant_set', 'templateinvariant_set')
-        variants = list(variants_query)
+        variants = list[Variant]()
+        for i in range(0, variants_query.count(), self.batch_size):
+            variants.extend(variants_query[i:i + self.batch_size])
+        variants_counts: dict[str, int] = {
+            i: c
+            for i, c in Variant
+            .objects
+            .order_by('pk')
+            .annotate(variant_count_updated=Count('of__variants', distinct=True, filter=Q(of__variants__status__in=Variant.public_statuses())))
+            .values_list('id', 'variant_count_updated')
+        }
         variants_to_save = update_variants(
             variants,
+            variants_counts,
             edhrec_variant_db,
             log=lambda x: self.log(x),
             log_warning=lambda x: self.log(x, self.style.WARNING),
             log_error=lambda x: self.log(x, self.style.ERROR),
         )
         updated_variant_count = len(variants_to_save)
-        Variant.objects.bulk_update(variants_to_save, fields=Variant.playable_fields() + ['popularity'], batch_size=5000)
+        Variant.objects.bulk_update(variants_to_save, fields=Variant.playable_fields() + ['popularity', 'variant_count'], batch_size=self.batch_size)
         self.log('Updating variants...done', self.style.SUCCESS)
         if updated_variant_count > 0:
             self.log(f'Successfully updated {updated_variant_count} variants', self.style.SUCCESS)

backend/spellbook/management/edhrec.py

@@ -39,7 +39,7 @@ def edhrec():
     return variants_db
 
 
-def update_variants(variants: list[Variant], edhrec: dict[str, dict], log=lambda t: print(t), log_warning=lambda t: print(t), log_error=lambda t: print(t)):
+def update_variants(variants: list[Variant], counts: dict[str, int], edhrec: dict[str, dict], log=lambda t: print(t), log_warning=lambda t: print(t), log_error=lambda t: print(t)):
     variants_to_save: list[Variant] = []
     for variant in variants:
         updated = False
@@ -57,6 +57,12 @@ def update_variants(variants: list[Variant], edhrec: dict[str, dict], log=lambda
             or any(tiv.must_be_commander for tiv in variant.templateinvariant_set.all())
         if variant.update(variant.uses.all(), requires_commander):
             updated = True
+        # Update with Spellbook data
+        variant_count = counts.get(variant.id, 0)
+        if variant.variant_count != variant_count:
+            variant.variant_count = variant_count
+            updated = True
+        # Save if updated
         if updated:
             variants_to_save.append(variant)
     return variants_to_save

backend/spellbook/management/scryfall.py

@@ -65,7 +65,7 @@ def fuzzy_restore_card(scryfall: dict, name: str):
             raise Exception(f'Card {name} not found in scryfall dataset, even after fuzzy search')
 
 
-def update_cards(cards: list[Card], scryfall: dict[str, dict], log=lambda t: print(t), log_warning=lambda t: print(t), log_error=lambda t: print(t)):
+def update_cards(cards: list[Card], scryfall: dict[str, dict], counts: dict[int, int], log=lambda t: print(t), log_warning=lambda t: print(t), log_error=lambda t: print(t)):
     oracle_db = {card_object['oracle_id']: card_object for card_object in scryfall.values()}
     existing_names = {card.name: card for card in cards}
     existing_oracle_ids = {card.oracle_id: card for card in cards if card.oracle_id is not None}
@@ -166,6 +166,10 @@ def card_fields(card: Card) -> tuple:
             log_warning(f'Card {card.name} with oracle id {oracle_id} not found in scryfall dataset. Oracle id has been removed.')
             card.oracle_id = None
             updated = True
+        variant_count = counts.get(card.id, 0)
+        if card.variant_count != variant_count:
+            card.variant_count = variant_count
+            updated = True
         if updated:
             cards_to_save.append(card)
     return cards_to_save

backend/spellbook/migrations/0035_variant_variant_count.py

@@ -0,0 +1,30 @@
+# Generated by Django 5.1.1 on 2024-10-18 23:17
+
+from django.db import migrations, models
+
+
+def init_variant_counts(apps, schema_editor):
+    Variant = apps.get_model('spellbook', 'Variant')
+    variants = list(Variant.objects.only('id').annotate(
+        variant_count_updated=models.Count('of__variants', distinct=True, filter=models.Q(of__variants__status__in=('OK', 'E')))
+    ))
+    for variant in variants:
+        variant.variant_count = variant.variant_count_updated
+        variant.pre_save = lambda: None
+    Variant.objects.bulk_update(variants, ['variant_count'])
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('spellbook', '0034_combo_card_count_combo_ingredient_count_and_more'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='variant',
+            name='variant_count',
+            field=models.PositiveIntegerField(default=0, editable=False, help_text='Number of variants generated by the same generator combos'),
+        ),
+        migrations.RunPython(init_variant_counts, reverse_code=migrations.RunPython.noop),
+    ]

backend/spellbook/models/variant.py

@@ -105,6 +105,7 @@ def preview_statuses(cls):
     description_line_count = models.PositiveIntegerField(editable=False, help_text='Number of lines in the description')
     other_prerequisites_line_count = models.PositiveIntegerField(editable=False, help_text='Number of lines in the other prerequisites')
     published = models.BooleanField(editable=False, default=False, help_text='Whether the variant has been published')
+    variant_count = models.PositiveIntegerField(editable=False, default=0, help_text='Number of variants generated by the same generator combos')
 
     class Meta:
         verbose_name = 'variant'

backend/spellbook/serializers/combo_serializer.py

@@ -7,5 +7,4 @@ class Meta:
         model = Combo
         fields = [
             'id',
-            'variant_count',
         ]

backend/spellbook/serializers/variant_serializer.py

@@ -167,6 +167,7 @@ class Meta:
             'spoiler',
             'legalities',
             'prices',
+            'variant_count',
         ]
 
     @classmethod

backend/spellbook/tests/test_views/test_variant_views.py

@@ -18,10 +18,11 @@ def setUp(self) -> None:
         super().generate_variants()
         Variant.objects.update(status=Variant.Status.OK)
         Variant.objects.filter(id__in=random.sample(list(Variant.objects.values_list('id', flat=True)), 3)).update(status=Variant.Status.EXAMPLE)
-        self.bulk_serialize_variants()
         self.v1_id: int = Variant.objects.first().id  # type: ignore
         self.public_variants = VariantViewSet.queryset
         self.ok_variants = self.public_variants.filter(status=Variant.Status.OK)
+        self.update_variants_count()
+        self.bulk_serialize_variants()
 
     def variant_assertions(self, variant_result):
         v: Variant = Variant.objects.get(id=variant_result.id)
@@ -114,6 +115,7 @@ def variant_assertions(self, variant_result):
         includes_list = [i.id for i in v.includes.all()]
         for i in variant_result.includes:
             self.assertIn(i.id, includes_list)
+        self.assertEqual(variant_result.variant_count, self.public_variants.filter(of__variants=v.id).values('id').distinct().count())
 
     def test_variants_list_view(self):
         response = self.client.get('/variants', follow=True)
@@ -790,24 +792,13 @@ def test_variants_list_view_grouping_by_combo(self):
             result_id_set = {v.id for v in result.results}
             self.assertSetEqual(result_id_set, best_variants_ids)
 
-    def test_variants_list_view_of_filter(self):
-        for combo_id in Combo.objects.values_list('pk', flat=True):
-            with self.subTest(f'combo {combo_id}'):
-                response = self.client.get('/variants', query_params={'of': combo_id}, follow=True)  # type: ignore
+    def test_variants_list_view_variant_filter(self):
+        for variant_id in Variant.objects.values_list('pk', flat=True):
+            with self.subTest(f'combo {variant_id}'):
+                response = self.client.get('/variants', query_params={'variant': variant_id}, follow=True)  # type: ignore
                 self.assertEqual(response.status_code, 200, response.content.decode())
                 self.assertEqual(response.get('Content-Type'), 'application/json')
                 result = json.loads(response.content, object_hook=json_to_python_lambda)
                 result_id_set = {v.id for v in result.results}
-                correct_id_set = {v.id for v in Variant.objects.filter(of=combo_id)}
-                self.assertSetEqual(result_id_set, correct_id_set)
-
-    def test_variants_list_view_includes_filter(self):
-        for combo_id in Combo.objects.values_list('pk', flat=True):
-            with self.subTest(f'combo {combo_id}'):
-                response = self.client.get('/variants', query_params={'includes': combo_id}, follow=True)  # type: ignore
-                self.assertEqual(response.status_code, 200)
-                self.assertEqual(response.get('Content-Type'), 'application/json')
-                result = json.loads(response.content, object_hook=json_to_python_lambda)
-                result_id_set = {v.id for v in result.results}
-                correct_id_set = {v.id for v in Variant.objects.filter(includes=combo_id)}
+                correct_id_set = {v.id for v in Variant.objects.filter(of__variants=variant_id)}
                 self.assertSetEqual(result_id_set, correct_id_set)

backend/spellbook/tests/testing.py

@@ -3,7 +3,7 @@
 from functools import reduce
 from collections import defaultdict
 
-from django.db.models import Count, OuterRef, Subquery
+from django.db.models import Q, Count, OuterRef, Subquery
 from django.db.models.functions import Coalesce
 from common.testing import TestCaseMixin as BaseTestCaseMixin
 from django.conf import settings
@@ -34,14 +34,34 @@ def update_variants_count(self):
                 Subquery(
                     Variant
                     .objects
-                    .filter(uses=OuterRef('pk'), status__in=('OK', 'E'))
+                    .filter(uses=OuterRef('pk'), status__in=Variant.public_statuses())
                     .values('uses')
-                    .annotate(total=Count('pk'))
+                    .annotate(total=Count('pk', distinct=True))
                     .values('total'),
                 ),
                 0,
             ),
         )
+        Combo.objects.update(
+            variant_count=Coalesce(
+                Subquery(
+                    Variant
+                    .objects
+                    .filter(of=OuterRef('pk'), status__in=Variant.public_statuses())
+                    .values('of')
+                    .annotate(total=Count('pk', distinct=True))
+                    .values('total'),
+                ),
+                0,
+            ),
+        )
+        variants = list(Variant.objects.only('id').annotate(
+            variant_count_updated=Count('of__variants', distinct=True, filter=Q(of__variants__status__in=Variant.public_statuses()))
+        ))
+        for variant in variants:
+            variant.variant_count = variant.variant_count_updated
+            variant.pre_save = lambda: None
+        Variant.objects.bulk_update(variants, ['variant_count'])
 
     def generate_and_publish_variants(self):
         self.generate_variants()

backend/spellbook/views/variants.py

@@ -4,15 +4,15 @@
 from django.template import loader
 from rest_framework import viewsets, serializers, filters
 from django_filters.rest_framework import DjangoFilterBackend, FilterSet
-from django_filters.filters import NumberFilter
+from django_filters.filters import CharFilter
 from drf_spectacular.utils import extend_schema, inline_serializer
 from spellbook.models import Variant, PreSerializedSerializer
 from spellbook.serializers import VariantSerializer
 from .filters import SpellbookQueryFilter, OrderingFilterWithNullsLast
 
 
 class VariantGroupedByComboFilter(filters.BaseFilterBackend):
-    DEFAULT_GROUPING_ORDERING = (F('popularity').desc(), F('identity_count').asc(), F('card_count').asc(), F('id').asc())
+    DEFAULT_GROUPING_ORDERING = (F('popularity').desc(), F('identity_count').asc(), F('card_count').asc())
     query_param = 'group_by_combo'
     template = 'spellbook/filters/group_by_combo.html'
 
@@ -30,7 +30,7 @@ def _filter_queryset(self, queryset: QuerySet[Variant]) -> QuerySet[Variant]:
             rank=Window(
                 expression=Rank(),
                 partition_by=F('variantofcombo__combo_id'),
-                order_by=queryset.query.order_by or self.DEFAULT_GROUPING_ORDERING,
+                order_by=(queryset.query.order_by or self.DEFAULT_GROUPING_ORDERING) + (F('pk'),),
             )
         ).filter(rank=1).values('id').distinct()
         return queryset.filter(id__in=top_variants_for_each_combo)
@@ -63,8 +63,7 @@ def to_html(self, request, queryset, view):
 
 
 class VariantFilterSet(FilterSet):
-    of = NumberFilter(field_name='of__pk', label='Variants generated by combo.')
-    includes = NumberFilter(field_name='includes__pk', label='Variants that include combo.')
+    variant = CharFilter(field_name='of__variants', label='Filters for variants of the same combos that generated the given variant id.', distinct=True)
 
 
 @extend_schema(responses={
@@ -92,6 +91,7 @@ class VariantViewSet(viewsets.ReadOnlyModelViewSet):
         'mana_value_needed',
         'description_line_count',
         'other_prerequisites_line_count',
+        'variant_count',
         'created',
         'updated',
         '?'