update WI group importer to handle councils

mysociety · Mar 26, 2024 · 834888a · 834888a
1 parent 490eeb3
commit 834888a
Showing 1 changed file with 15 additions and 3 deletions.
diff --git a/hub/management/commands/import_wi_group_locations.py b/hub/management/commands/import_wi_group_locations.py
@@ -6,15 +6,27 @@
 
 from hub.models import DataSet
 
-from .base_importers import BaseConstituencyGroupListImportCommand
+from .base_importers import (
+    BaseConstituencyGroupListImportCommand,
+    MultipleAreaTypesMixin,
+)
 
 
-class Command(BaseConstituencyGroupListImportCommand):
+class Command(MultipleAreaTypesMixin, BaseConstituencyGroupListImportCommand):
     help = "Import data about WI groups per constituency"
     message = "Importing Women's Institute group data"
 
     data_file = settings.BASE_DIR / "data" / "wi_groups.csv"
     source_url = "https://www.thewi.org.uk/wis-a-z"
+
+    uses_gss = True
+    area_types = ["WMC", "WMC23", "STC", "DIS"]
+    cons_col_map = {
+        "WMC": "WMC",
+        "WMC23": "WMC23",
+        "STC": "STC",
+        "DIS": "DIS",
+    }
     defaults = {
         "label": "Women’s Institute groups",
         "data_type": "json",
@@ -70,7 +82,7 @@ def get_df(self):
         df.group_name = df.group_name.apply(
             lambda x: x.split(" | ")[0] if isinstance(x, str) else x
         )
-        df.columns = ["group_name", "url", "lat_lon", "constituency"]
+        df.columns = ["group_name", "url", "lat_lon", "constituency", *self.area_types]
         return df
 
     def get_group_json(self, row):