From 30fb05bc60cec015a7514c8e3cbebec90647a14b Mon Sep 17 00:00:00 2001
From: Struan Donald <struan@exo.org.uk>
Date: Wed, 24 Apr 2024 12:08:01 +0100
Subject: [PATCH] fix constituency country importer for multiple areas

---
 .../commands/import_area_countries.py         | 101 +++++++++---------
 1 file changed, 52 insertions(+), 49 deletions(-)

diff --git a/hub/management/commands/import_area_countries.py b/hub/management/commands/import_area_countries.py
index 37c7e7389..63a0d2e6d 100644
--- a/hub/management/commands/import_area_countries.py
+++ b/hub/management/commands/import_area_countries.py
@@ -1,68 +1,71 @@
 from datetime import date
 
-from django.core.management.base import BaseCommand
+import pandas as pd
 
-from tqdm import tqdm
-
-from hub.models import Area, AreaData, DataSet, DataType
+from hub.models import DataSet
 from utils.mapit import MapIt
 
+from .base_importers import BaseImportFromDataFrameCommand, MultipleAreaTypesMixin
+
 
-class Command(BaseCommand):
+class Command(MultipleAreaTypesMixin, BaseImportFromDataFrameCommand):
     help = "Import countries of areas from MapIt"
+    message = "Importing constituency countries"
+    cons_row = "gss-code"
+    uses_gss = True
+    do_not_convert = True
 
     def add_arguments(self, parser):
         parser.add_argument(
             "-q", "--quiet", action="store_true", help="Silence progress bars."
         )
 
-    def handle(self, quiet=False, *args, **options):
-        self._quiet = quiet
-        data_type = self.create_data_type()
-        self.delete_data(data_type)
-        self.import_data(data_type)
+    area_types = ["WMC", "WMC23"]
 
-    def create_data_type(self):
-        options = [
-            {"title": "England", "shader": "#f8f9fa"},
-            {"title": "Wales", "shader": "#cc3517"},
-            {"title": "Scotland", "shader": "#202448"},
-            {"title": "Northern Ireland", "shader": "#458945"},
-        ]
-        countries_ds, created = DataSet.objects.update_or_create(
-            name="country",
-            defaults={
-                "data_type": "text",
-                "description": "",
-                "release_date": str(date.today()),
-                "label": "Country of the UK",
-                "source_label": "Data from mySociety.",
-                "source": "https://mapit.mysociety.org/",
-                "table": "areadata",
-                "options": options,
-                "comparators": DataSet.in_comparators(),
-            },
-        )
+    mapit_types = {
+        "WMC": ["WMC"],
+        "WMC23": ["WMCF"],
+    }
 
-        countries, created = DataType.objects.update_or_create(
-            data_set=countries_ds,
-            name="country",
-            defaults={"data_type": "text"},
-        )
+    options = [
+        {"title": "England", "shader": "#f8f9fa"},
+        {"title": "Wales", "shader": "#cc3517"},
+        {"title": "Scotland", "shader": "#202448"},
+        {"title": "Northern Ireland", "shader": "#458945"},
+    ]
 
-        return countries
+    defaults = {
+        "data_type": "text",
+        "category": "place",
+        "description": "",
+        "release_date": str(date.today()),
+        "label": "Country of the UK",
+        "source_label": "Data from mySociety.",
+        "source": "https://mapit.mysociety.org/",
+        "table": "areadata",
+        "options": options,
+        "comparators": DataSet.in_comparators(),
+        "unit_type": "raw",
+        "unit_distribution": "",
+        "is_shadable": True,
+        "is_filterable": True,
+        "is_public": True,
+    }
 
-    def import_data(self, data_type):
-        mapit = MapIt()
-        areas = mapit.areas_of_type(
-            [area.mapit_id for area in Area.objects.exclude(mapit_id="")]
-        )
-        for area in tqdm(areas, disable=self._quiet):
-            data, created = AreaData.objects.update_or_create(
-                area=Area.objects.get(mapit_id=area["id"]),
-                data_type=data_type,
-                data=area["country_name"],
+    data_sets = {"country": {"defaults": defaults, "col": "country"}}
+
+    def get_dataframe(self):
+        mapit_client = MapIt()
+        areas = mapit_client.areas_of_type(self.mapit_types[self.area_type])
+        types = []
+        for area in areas:
+            types.append(
+                {
+                    "gss-code": area["codes"]["gss"],
+                    "country": area["country_name"],
+                }
             )
 
-    def delete_data(self, data_type):
-        AreaData.objects.filter(data_type=data_type).delete()
+        df = pd.DataFrame(types)
+
+        return df