update election imports to set area types

mysociety · Jul 10, 2024 · 3a907ce · 3a907ce
1 parent 932089b
commit 3a907ce
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 7 deletions.
diff --git a/hub/management/commands/import_last_election_data.py b/hub/management/commands/import_last_election_data.py
@@ -235,6 +235,7 @@ def create_data_types(self):
         second_party, created = DataType.objects.update_or_create(
             data_set=second_party_ds,
             name="second_party",
+            area_type=self.get_area_type(),
             defaults={"data_type": "text"},
         )
 
@@ -258,6 +259,7 @@ def create_data_types(self):
         last_election, created = DataType.objects.update_or_create(
             data_set=last_election_ds,
             name="last_election",
+            area_type=self.get_area_type(),
             defaults={"data_type": "json"},
         )
 

diff --git a/hub/management/commands/import_mps_election_results.py b/hub/management/commands/import_mps_election_results.py
@@ -6,19 +6,29 @@
 from dateutil.parser import isoparse
 from tqdm import tqdm
 
-from hub.models import DataSet, DataType, Person, PersonData
+from hub.models import AreaType, DataSet, DataType, Person, PersonData
 
 
 class Command(BaseCommand):
     help = "Import election results for UK Members of Parliament"
 
+    area_type = "WMC23"
+
     def add_arguments(self, parser):
         parser.add_argument(
             "-q", "--quiet", action="store_true", help="Silence progress bars."
         )
 
+        parser.add_argument(
+            "--area_type", action="store", help="Set area type code, default is WMC23"
+        )
+
     def handle(self, quiet=False, *args, **options):
         self._quiet = quiet
+
+        if options.get("area_type") is not None:
+            self.area_type = options["area_type"]
+
         self.import_results()
 
     def get_results(self):
@@ -48,6 +58,9 @@ def get_results(self):
                     mp_id.person.end_date = end_date.date().isoformat()
                     mp_id.person.save()
                     continue
+                else:
+                    mp_id.person.end_date = None
+                    mp_id.person.save()
 
                 results[mp_id.person.id] = {
                     "first_elected": membership["membershipStartDate"]
@@ -76,6 +89,7 @@ def get_results(self):
     def create_data_types(self):
         if not self._quiet:
             self.stdout.write("Creating data sets and types")
+        area_type = AreaType.objects.get(code=self.area_type)
         majority_ds, created = DataSet.objects.update_or_create(
             name="mp_election_majority",
             defaults={
@@ -90,9 +104,11 @@ def create_data_types(self):
                 "default_value": 1000,
             },
         )
+        majority_ds.areas_available.add(area_type)
         majority, created = DataType.objects.update_or_create(
             data_set=majority_ds,
             name="mp_election_majority",
+            area_type=area_type,
             defaults={"label": "MP majority", "data_type": "integer"},
         )
 
@@ -110,6 +126,7 @@ def create_data_types(self):
                 "default_value": 2024,
             },
         )
+        last_elected_ds.areas_available.add(area_type)
         last_elected, created = DataType.objects.update_or_create(
             data_set=last_elected_ds,
             name="mp_last_elected",
@@ -130,6 +147,7 @@ def create_data_types(self):
                 "default_value": 2024,
             },
         )
+        first_elected_ds.areas_available.add(area_type)
         first_elected, created = DataType.objects.update_or_create(
             data_set=first_elected_ds,
             name="mp_first_elected",
@@ -149,11 +167,12 @@ def add_results(self, results, data_types):
             person = Person.objects.get(id=mp_id)
 
             for key, data_type in data_types.items():
-                data, created = PersonData.objects.update_or_create(
-                    person=person,
-                    data_type=data_type,
-                    defaults={"data": result[key]},
-                )
+                if result.get(key) is not None:
+                    data, created = PersonData.objects.update_or_create(
+                        person=person,
+                        data_type=data_type,
+                        defaults={"data": result[key]},
+                    )
 
     def import_results(self):
         data_types = self.create_data_types()

diff --git a/hub/tests/test_import_mps.py b/hub/tests/test_import_mps.py
@@ -60,7 +60,7 @@ def test_import(self, id_df_mock, twfy_df_mock, retrieve_mock, request_mock):
                     26,
                     27,
                 ],
-                "Constituency": ["South Borsetshire", "Borsetshire West"],
+                "Constituency": ["New South Borsetshire", "New Borsetshire West"],
                 "Party": ["Borsetshire Independence", "Borsetshire Unionist"],
                 "First name": ["James", "Angela"],
                 "Last name": ["Madeupname", "Madeupname"],