Generalized aircraft assign function

cntmanage/flightslot/admins/mission_adm.py

@@ -1,7 +1,7 @@
-from django.forms import ModelMultipleChoiceField
+from django.forms import TypedMultipleChoiceField
 from django.contrib import admin, messages
 from django.http import HttpRequest
-from django.db.models.query import QuerySet, Q
+from django.db.models.query import QuerySet
 from django.utils.safestring import SafeText
 
 from django_admin_action_forms import AdminActionFormsMixin, AdminActionForm, action_with_form
@@ -11,11 +11,13 @@ from import_export.admin import ImportMixin
 from import_export.tmp_storages import CacheStorage
 from import_export.resources import ModelResource
 
-from ..models.aircrafts import Aircraft
+from ..models.aircrafts import AircraftTypes
 from ..models.missions import MissionProfile
 
+from ..actions.assign_aircraft import assign_aircraft
+
 from datetime import timedelta
-from typing import Any, Dict
+from typing import Any, Dict, List
 
 # Resource Class for Student data import
 class MissionProfileResource(ModelResource):
@@ -40,7 +42,7 @@ class MissionProfileResource(ModelResource):
 
 # Form class to assing aircrafts to students
 class ChangeAircraftForm(AdminActionForm):
-    aircrafts = ModelMultipleChoiceField(queryset=Aircraft.objects.distinct('type').all())
+    aircrafts = TypedMultipleChoiceField(choices=AircraftTypes)
 
 class MissionProfileAdmin(ImportMixin, AdminActionFormsMixin, admin.ModelAdmin):
     list_display = ("mtype", "mnum", "assigned_aircrafts", "duration", "notes", )
@@ -53,20 +55,11 @@ class MissionProfileAdmin(ImportMixin, AdminActionFormsMixin, admin.ModelAdmin):
     def get_queryset(self, request: HttpRequest) -> QuerySet[MissionProfile]:
         return super().get_queryset(request).order_by("mtype", "mnum")
 
-    @action_with_form(ChangeAircraftForm, description="Assign Aircraft")
-    def assign_aircraft(self, request: HttpRequest, queryset: QuerySet[MissionProfile], data: Dict[str, QuerySet[Aircraft]]):
-        ac_types = [t.type for t in data["aircrafts"]]
-        ac_query: Q = Q() # Build an or query to select all aircrafts of the specified types
-        for a in ac_types:
-            ac_query |= Q(type=a)
-        aircrafts: QuerySet[Aircraft] = Aircraft.objects.filter(ac_query).all() # Execute query
-        i: int = 0
-        for mix in queryset:
-            mix.aircrafts.clear()
-            for ac in aircrafts:
-                mix.aircrafts.add(ac)
-            mix.save()
-            i += 1
+    @action_with_form(ChangeAircraftForm, description="Assign Aircraft Type")
+    def assign_aircraft(self, request: HttpRequest, queryset: QuerySet[MissionProfile], data: Dict[str, List[AircraftTypes]]):
+        i: int
+        ac_types: List[str]
+        i, ac_types = assign_aircraft(queryset=queryset, data=data)
         messages.success(request, f"{i} Missions updated to {ac_types}")
 
     @admin.display(description="Assigned Aircrafts")