Refactor country filter to be cleaner (#96)

IGES-Geospatial · Feb 9, 2022 · a6515bd · a6515bd
1 parent 282ba7d
commit a6515bd
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/go_utils/geoenrich.py b/go_utils/geoenrich.py
@@ -85,17 +85,14 @@ def get_country_api_data(
     if is_clean:
         df = default_data_clean(df, protocol)
 
+    for region in regions:
+        countries.extend(region_dict[region])
+    countries_set = set(countries)
     # Return the regular data if nothing is specified
-    if not countries and not regions:
+    if not countries_set:
         return df
     else:
-        mask = np.zeros((len(df)), dtype=bool)
-        if countries:
-            mask |= _get_valid_countries_mask(df, protocol, countries)
-
-        if regions:
-            for region in regions:
-                mask |= _get_valid_countries_mask(df, protocol, region_dict[region])
+        mask = _get_valid_countries_mask(df, protocol, countries_set)
         return df[mask]