Sprinkle of typing and import juggling to facilitate

polling_stations/apps/councils/models.py

@@ -1,5 +1,6 @@
 import re
 from pathlib import Path
+from typing import Optional
 
 from data_importers.event_types import DataEventType
 from data_importers.models import DataEvent, DataQuality
@@ -11,6 +12,7 @@
 from django.db.models import Count, JSONField, OuterRef, Q, Subquery
 from django.utils.translation import get_language
 from file_uploads.models import File, Upload
+from pollingstations.models import PollingStation
 
 from polling_stations.i18n.cy import WelshNameMutationMixin
 
@@ -180,7 +182,7 @@ def has_polling_stations_in_db(self):
             return True
         return False
 
-    def latest_data_event(self, event_type):
+    def latest_data_event(self, event_type: DataEventType):
         data_event_model = apps.get_model("data_importers", "DataEvent")
         try:
             return self.dataevent_set.filter(event_type=event_type).latest()
@@ -204,11 +206,13 @@ def live_upload(self):
             return latest_import.upload
         return None
 
-    def update_all_station_visibilities_from_events(self, election_dates):
+    def update_all_station_visibilities_from_events(self, election_dates: list):
         for station in self.pollingstation_set.all():
             self.update_station_visibility_from_events(station, election_dates)
 
-    def update_station_visibility_from_events(self, station, election_dates=None):
+    def update_station_visibility_from_events(
+        self, station: PollingStation, election_dates: Optional[list] = None
+    ):
         query = Q(
             council_id=self.council_id,
             event_type=DataEventType.SET_STATION_VISIBILITY,

polling_stations/apps/data_finder/helpers/geocoders.py

@@ -1,8 +1,8 @@
 import abc
 
+from councils.models import Council
 from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
-from pollingstations.models import Council
 from uk_geo_utils.geocoders import (
     AddressBaseGeocoder,
     CodesNotFoundException,

polling_stations/apps/data_importers/event_helpers.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from councils.models import Council
 from data_importers.event_types import DataEventType
 from data_importers.models import DataEvent
@@ -8,14 +10,18 @@ class EventValidationError(Exception):
     pass
 
 
-def record_teardown_event(council_id):
+def record_teardown_event(council_id: str):
     DataEvent.objects.create(
         council=Council.objects.get(council_id=council_id),
         event_type=DataEventType.TEARDOWN,
     )
 
 
-def record_set_station_visibility_event(station, visibility_choice, metadata=None):
+def record_set_station_visibility_event(
+    station: PollingStation,
+    visibility_choice: VisibilityChoices,
+    metadata: Optional[dict] = None,
+):
     if metadata is None:
         metadata = {}
     council = station.council
@@ -36,7 +42,7 @@ def record_set_station_visibility_event(station, visibility_choice, metadata=Non
     )
 
 
-def set_station_visibility(event):
+def set_station_visibility(event: DataEvent):
     payload = event.payload
     if not (visibility := payload.get("visibility", None)):
         raise EventValidationError(f"{event} payload missing visibility.")

polling_stations/apps/pollingstations/models.py

@@ -6,7 +6,6 @@
 from itertools import groupby
 
 from core.opening_times import OpeningTimes
-from councils.models import Council
 from django.contrib.gis.db import models
 from django.core.exceptions import ObjectDoesNotExist
 from django.db.models import JSONField
@@ -17,7 +16,7 @@
 
 class PollingDistrict(models.Model):
     name = models.CharField(blank=True, null=True, max_length=255)
-    council = models.ForeignKey(Council, null=True, on_delete=models.CASCADE)
+    council = models.ForeignKey("councils.Council", null=True, on_delete=models.CASCADE)
     internal_council_id = models.CharField(blank=True, max_length=100)
     extra_id = models.CharField(blank=True, null=True, max_length=100)
     area = models.MultiPolygonField(null=True, blank=True)
@@ -44,7 +43,7 @@ class VisibilityChoices(models.TextChoices):
 
 class PollingStation(models.Model):
     council = models.ForeignKey(
-        Council, null=True, db_index=True, on_delete=models.CASCADE
+        "councils.Council", null=True, db_index=True, on_delete=models.CASCADE
     )
     internal_council_id = models.CharField(blank=True, max_length=100, db_index=True)
     postcode = models.CharField(blank=True, null=True, max_length=100)
@@ -167,7 +166,7 @@ class AdvanceVotingStation(models.Model):
     address = models.TextField(blank=True, null=True)
     location = models.PointField(null=True, blank=True)
     opening_times = JSONField(null=True)
-    council = models.ForeignKey(Council, null=True, on_delete=models.CASCADE)
+    council = models.ForeignKey("councils.Council", null=True, on_delete=models.CASCADE)
 
     def __str__(self):
         return f"{self.name} ({self.postcode})"