From beed784f7cbd9ad50c0330e908731e7a6ad65883 Mon Sep 17 00:00:00 2001
From: David Mandelberg <david@mandelberg.org>
Date: Sun, 5 Nov 2023 21:47:20 -0500
Subject: [PATCH] Add a field for additional Wikidata items in a media item

---
 rock_paper_sand/media_item.py      | 18 ++++++++++++++++--
 rock_paper_sand/media_item_test.py |  9 +++++++++
 rock_paper_sand/proto/config.proto |  9 +++++++--
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/rock_paper_sand/media_item.py b/rock_paper_sand/media_item.py
index 561adbe..005f515 100644
--- a/rock_paper_sand/media_item.py
+++ b/rock_paper_sand/media_item.py
@@ -42,7 +42,10 @@ class MediaItem:
         custom_data: Parsed proto.custom_data field, or None.
         done: Parsed proto.done field.
         wikidata_item: Wikidata item, or None.
-        all_wikidata_items_recursive: Wikidata items from this and all parts.
+        all_wikidata_items: wikidata_item and also any items in the
+            proto.wikidata_additional field.
+        all_wikidata_items_recursive: Wikidata items from this and all parts,
+            including ones from wikidata_additional.
         wikidata_ignore_items_recursive: Wikidata items to ignore from this and
             all parts.
         wikidata_classes_ignore_recursive: Wikidata classes to ignore from this
@@ -62,6 +65,7 @@ class MediaItem:
     custom_data: Any
     done: multi_level_set.MultiLevelSet
     wikidata_item: wikidata_value.Item | None
+    all_wikidata_items: Set[wikidata_value.Item]
     all_wikidata_items_recursive: Set[wikidata_value.Item]
     wikidata_ignore_items_recursive: Set[wikidata_value.Item]
     wikidata_classes_ignore_recursive: Set[wikidata_value.Item]
@@ -117,9 +121,18 @@ def from_config(
                 if proto.wikidata
                 else None
             )
-            all_wikidata_items_recursive = frozenset(
+            all_wikidata_items = frozenset(
                 itertools.chain(
                     () if wikidata_item is None else (wikidata_item,),
+                    map(
+                        wikidata_value.Item.from_string,
+                        proto.wikidata_additional,
+                    ),
+                )
+            )
+            all_wikidata_items_recursive = frozenset(
+                itertools.chain(
+                    all_wikidata_items,
                     *(part.all_wikidata_items_recursive for part in parts),
                 )
             )
@@ -160,6 +173,7 @@ def from_config(
                 ),
                 done=multi_level_set.MultiLevelSet.from_string(proto.done),
                 wikidata_item=wikidata_item,
+                all_wikidata_items=all_wikidata_items,
                 all_wikidata_items_recursive=all_wikidata_items_recursive,
                 wikidata_ignore_items_recursive=wikidata_ignore_items_recursive,
                 wikidata_classes_ignore_recursive=(
diff --git a/rock_paper_sand/media_item_test.py b/rock_paper_sand/media_item_test.py
index 78259b9..71dc54b 100644
--- a/rock_paper_sand/media_item_test.py
+++ b/rock_paper_sand/media_item_test.py
@@ -32,6 +32,7 @@ def test_from_config(self) -> None:
         other_part = {
             "name": "other-part",
             "wikidata": "Q2",
+            "wikidataAdditional": ["Q20"],
             "wikidataIgnore": ["Q21", "Q22"],
             "wikidataClassesIgnore": ["Q23", "Q24"],
         }
@@ -65,9 +66,11 @@ def test_from_config(self) -> None:
                 custom_data={"a": "b"},
                 done=mock.ANY,
                 wikidata_item=wikidata_value.Item("Q1"),
+                all_wikidata_items={wikidata_value.Item("Q1")},
                 all_wikidata_items_recursive={
                     wikidata_value.Item("Q1"),
                     wikidata_value.Item("Q2"),
+                    wikidata_value.Item("Q20"),
                 },
                 wikidata_ignore_items_recursive={
                     wikidata_value.Item("Q11"),
@@ -91,6 +94,7 @@ def test_from_config(self) -> None:
                         custom_data=None,
                         done=mock.ANY,
                         wikidata_item=None,
+                        all_wikidata_items=frozenset(),
                         all_wikidata_items_recursive=frozenset(),
                         wikidata_ignore_items_recursive=frozenset(),
                         wikidata_classes_ignore_recursive=frozenset(),
@@ -107,8 +111,13 @@ def test_from_config(self) -> None:
                         custom_data=None,
                         done=mock.ANY,
                         wikidata_item=wikidata_value.Item("Q2"),
+                        all_wikidata_items={
+                            wikidata_value.Item("Q2"),
+                            wikidata_value.Item("Q20"),
+                        },
                         all_wikidata_items_recursive={
                             wikidata_value.Item("Q2"),
+                            wikidata_value.Item("Q20"),
                         },
                         wikidata_ignore_items_recursive={
                             wikidata_value.Item("Q21"),
diff --git a/rock_paper_sand/proto/config.proto b/rock_paper_sand/proto/config.proto
index 26e7999..7b20a0b 100644
--- a/rock_paper_sand/proto/config.proto
+++ b/rock_paper_sand/proto/config.proto
@@ -300,10 +300,15 @@ message MediaItem {
   // URL.
   string custom_availability = 4;
 
-  // Wikidata item URL or QID, e.g., "https://www.wikidata.org/wiki/Q3107329" or
-  // "Q3107329".
+  // Primary Wikidata item URL or QID, e.g.,
+  // "https://www.wikidata.org/wiki/Q3107329" or "Q3107329".
   string wikidata = 9;
 
+  // Additional Wikidata items. E.g., a media item with parts might have both a
+  // Wikidata item for a fictional universe and a Wikidata item for a media
+  // franchise.
+  repeated string wikidata_additional = 13;
+
   // Wikidata items to ignore when checking for related media.
   repeated string wikidata_ignore = 11;