Add support for merge migrations in count aggregation

gwax · Jan 5, 2023 · 681fc1f · 681fc1f
1 parent 63160ac
commit 681fc1f
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 1 deletion.
diff --git a/mtg_ssm/containers/counts.py b/mtg_ssm/containers/counts.py
@@ -47,6 +47,8 @@ def aggregate_card_counts(
             if value:
                 counts[count_type] = value + counts.get(count_type, 0)
         if counts:
+            while scryfall_id in oracle.index.migrate_old_id_to_new_id:
+                scryfall_id = oracle.index.migrate_old_id_to_new_id[scryfall_id]
             if scryfall_id not in oracle.index.id_to_card:
                 raise CardNotFoundError(
                     f"Found counts for card={scryfall_id} not found scryfall data"

diff --git a/mtg_ssm/containers/indexes.py b/mtg_ssm/containers/indexes.py
@@ -7,7 +7,7 @@
 
 from mtg_ssm.containers.bundles import ScryfallDataSet
 from mtg_ssm.mtg import util
-from mtg_ssm.scryfall.models import ScryCard, ScrySet
+from mtg_ssm.scryfall.models import ScryCard, ScryMigrationStrategy, ScrySet
 
 
 def name_card_sort_key(card: ScryCard) -> Tuple[str, int, str]:
@@ -56,6 +56,7 @@ def __init__(self) -> None:
         self.setcode_to_cards: Dict[str, List[ScryCard]] = {}
         self.id_to_setindex: Dict[UUID, int] = {}
         self.setcode_to_set: Dict[str, ScrySet] = {}
+        self.migrate_old_id_to_new_id: Dict[UUID, UUID] = {}
         # snnma = Set, Name, (Collector) Number, Multiverse ID, Artist
         # TODO: convert to intersecting bitmap indexes
         # TODO: do we really need artist?
@@ -99,6 +100,15 @@ def load_data(self, scrydata: ScryfallDataSet) -> None:
             self.id_to_setindex.update({c.id: i for i, c in enumerate(cards_list)})
         self.setcode_to_cards = dict(setcode_to_unsorted_cards)
 
+        for migration in scrydata.migrations:
+            if (
+                migration.migration_strategy == ScryMigrationStrategy.MERGE
+                and migration.new_scryfall_id
+            ):
+                self.migrate_old_id_to_new_id[
+                    migration.old_scryfall_id
+                ] = migration.new_scryfall_id
+
 
 class Oracle:
     """Container for an indexed Scryfall data set."""

diff --git a/tests/containers/test_counts.py b/tests/containers/test_counts.py
@@ -181,6 +181,20 @@ def test_diff_card_counts(
             },
             id="duplicates",
         ),
+        pytest.param(
+            [
+                {
+                    "scryfall_id": UUID("585fa2cc-4f77-47ab-8d2c-c68258ced283"),
+                    "foil": 1,
+                },
+            ],
+            {
+                UUID("9052f5c7-ee3b-457d-97ca-ac6b4518997c"): {
+                    counts.CountType.FOIL: 1,
+                }
+            },
+            id="migration",
+        ),
     ],
 )
 def test_aggregate_card_counts(