From e7aead4063b6a6f1e3645b7b555f0eef6499995f Mon Sep 17 00:00:00 2001
From: april <aprl.petersen@gmail.com>
Date: Thu, 11 Jan 2024 16:40:19 -0600
Subject: [PATCH] Make totals by class use human-readable names

---
 api/database/flights.py | 16 ++++++++++++----
 api/schemas/aircraft.py |  3 +++
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/api/database/flights.py b/api/database/flights.py
index d711a5b..395a78c 100644
--- a/api/database/flights.py
+++ b/api/database/flights.py
@@ -6,7 +6,8 @@ from bson import ObjectId
 from bson.errors import InvalidId
 from fastapi import HTTPException
 
-from schemas.aircraft import AircraftCreateSchema, aircraft_add_helper, AircraftCategory, AircraftClass
+from schemas.aircraft import AircraftCreateSchema, aircraft_add_helper, AircraftCategory, AircraftClass, \
+    aircraft_class_dict
 from .aircraft import retrieve_aircraft_by_tail, update_aircraft, update_aircraft_field, retrieve_aircraft
 from .db import flight_collection, aircraft_collection
 from schemas.flight import FlightConciseSchema, FlightDisplaySchema, FlightCreateSchema, flight_display_helper, \
@@ -107,12 +108,19 @@ async def retrieve_totals(user: str, start_date: datetime = None, end_date: date
 
     cursor = aircraft_collection.aggregate(pipeline)
 
-    result = await cursor.to_list(None)
+    result_list = await cursor.to_list(None)
 
-    if not result:
+    if not result_list:
         return {}
 
-    return dict(result[0])
+    result = dict(result_list[0])
+
+    print(aircraft_class_dict)
+
+    for entry in result["by_class"]:
+        entry["aircraft_class"] = aircraft_class_dict[entry["aircraft_class"]]
+
+    return result
 
 
 async def retrieve_flight(id: str) -> FlightDisplaySchema:
diff --git a/api/schemas/aircraft.py b/api/schemas/aircraft.py
index 0fff7cc..cb219a7 100644
--- a/api/schemas/aircraft.py
+++ b/api/schemas/aircraft.py
@@ -78,6 +78,9 @@ class AircraftClass(Enum):
     wss = "Weight-Shift Control Sea"
 
 
+aircraft_class_dict = {cls.name: cls.value for cls in AircraftClass}
+
+
 class AircraftCreateSchema(BaseModel):
     tail_no: str
     make: str