fix: added country to search endpoint response

MobilityData · Aug 9, 2024 · f73fc30 · f73fc30
1 parent 799fd33
commit f73fc30
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/api/src/feeds/impl/models/search_feed_item_result_impl.py b/api/src/feeds/impl/models/search_feed_item_result_impl.py
@@ -1,6 +1,7 @@
 from feeds_gen.models.latest_dataset import LatestDataset
 from feeds_gen.models.search_feed_item_result import SearchFeedItemResult
 from feeds_gen.models.source_info import SourceInfo
+import pycountry
 
 
 class SearchFeedItemResultImpl(SearchFeedItemResult):
@@ -83,6 +84,8 @@ def _translate_locations(cls, feed_search_row):
 
         for location in feed_search_row.locations:
             location["country"] = country_translations.get(location["country"], location["country"])
+            if location["country"] is None:
+                location["country"] = pycountry.countries.get(alpha_2=location["country_code"]).name
             location["subdivision_name"] = subdivision_translations.get(
                 location["subdivision_name"], location["subdivision_name"]
             )