migrate Enum to Literal when possible

Author: GabrielBarberini

lib/controllers/environment.py

@@ -41,7 +41,7 @@ async def get_rocketpy_env_binary(
         return env_service.get_env_binary()
 
     @controller_exception_handler
-    async def simulate_env(
+    async def get_environment_simulation(
         self, env_id: str
     ) -> EnvironmentSummary:
         """

lib/models/environment.py

@@ -1,18 +1,8 @@
 import datetime
-from enum import Enum
-from typing import Optional, ClassVar, Self
+from typing import Optional, ClassVar, Self, Literal
 from lib.models.interface import ApiBaseModel
 
 
-class AtmosphericModelTypes(str, Enum):
-    STANDARD_ATMOSPHERE: str = "STANDARD_ATMOSPHERE"
-    CUSTOM_ATMOSPHERE: str = "CUSTOM_ATMOSPHERE"
-    WYOMING_SOUNDING: str = "WYOMING_SOUNDING"
-    FORECAST: str = "FORECAST"
-    REANALYSIS: str = "REANALYSIS"
-    ENSEMBLE: str = "ENSEMBLE"
-
-
 class EnvironmentModel(ApiBaseModel):
     NAME: ClassVar = 'environment'
     METHODS: ClassVar = ('POST', 'GET', 'PUT', 'DELETE')
@@ -21,9 +11,7 @@ class EnvironmentModel(ApiBaseModel):
     elevation: Optional[int] = 1
 
     # Optional parameters
-    atmospheric_model_type: AtmosphericModelTypes = (
-        AtmosphericModelTypes.STANDARD_ATMOSPHERE
-    )
+    atmospheric_model_type: Literal['standard_atmosphere', 'custom_atmosphere', 'wyoming_sounding', 'forecast', 'reanalysis', 'ensemble'] = 'standard_atmosphere'
     atmospheric_model_file: Optional[str] = None
     date: Optional[datetime.datetime] = (
         datetime.datetime.today() + datetime.timedelta(days=1)

lib/models/flight.py

@@ -1,15 +1,9 @@
-from enum import Enum
-from typing import Optional, Self, ClassVar
+from typing import Optional, Self, ClassVar, Literal
 from lib.models.interface import ApiBaseModel
 from lib.models.rocket import RocketModel
 from lib.models.environment import EnvironmentModel
 
 
-class EquationsOfMotion(str, Enum):
-    STANDARD: str = "STANDARD"
-    SOLID_PROPULSION: str = "SOLID_PROPULSION"
-
-
 class FlightModel(ApiBaseModel):
     NAME: ClassVar = "flight"
     METHODS: ClassVar = ("POST", "GET", "PUT", "DELETE")
@@ -20,7 +14,7 @@ class FlightModel(ApiBaseModel):
     rail_length: float = 1
     time_overshoot: bool = True
     terminate_on_apogee: bool = True
-    equations_of_motion: EquationsOfMotion = EquationsOfMotion.STANDARD
+    equations_of_motion: Literal['standard', 'solid_propulsion'] = 'standard'
 
     # Optional parameters
     inclination: Optional[int] = None

lib/models/motor.py

@@ -3,7 +3,7 @@
 from pydantic import PrivateAttr, model_validator
 
 from lib.models.interface import ApiBaseModel
-from lib.models.sub.tanks import MotorTank, InterpolationMethods, TankCoordinateSystemOrientation
+from lib.models.sub.tanks import MotorTank
 
 
 class MotorKinds(str, Enum):
@@ -48,10 +48,8 @@ class MotorModel(ApiBaseModel):
     throat_radius: Optional[float] = None
 
     # Optional parameters
-    interpolation_method: InterpolationMethods = InterpolationMethods.LINEAR
-    coordinate_system_orientation: TankCoordinateSystemOrientation = (
-        TankCoordinateSystemOrientation.NOZZLE_TO_COMBUSTION_CHAMBER
-    )
+    interpolation_method: Literal['linear', 'spline', 'akima', 'polynomial', 'shepard', 'rbf'] = 'linear'
+    coordinate_system_orientation: Literal['nozzle_to_combustion_chamber', 'combustion_chamber_to_nozzle'] = 'nose_to_combustion_chamber'
     reshape_thrust_curve: Union[bool, tuple] = False
 
     # Computed parameters

lib/models/rocket.py

@@ -1,5 +1,4 @@
-from enum import Enum
-from typing import Optional, Tuple, List, Union, Self, ClassVar
+from typing import Optional, Tuple, List, Union, Self, ClassVar, Literal
 from lib.models.interface import ApiBaseModel
 from lib.models.motor import MotorModel
 from lib.models.sub.aerosurfaces import (
@@ -11,11 +10,6 @@
 )
 
 
-class RocketCoordinateSystemOrientation(str, Enum):
-    TAIL_TO_NOSE: str = "TAIL_TO_NOSE"
-    NOSE_TO_TAIL: str = "NOSE_TO_TAIL"
-
-
 class RocketModel(ApiBaseModel):
     NAME: ClassVar = "rocket"
     METHODS: ClassVar = ("POST", "GET", "PUT", "DELETE")
@@ -32,9 +26,7 @@ class RocketModel(ApiBaseModel):
     ] = (0, 0, 0)
     power_off_drag: List[Tuple[float, float]] = [(0, 0)]
     power_on_drag: List[Tuple[float, float]] = [(0, 0)]
-    coordinate_system_orientation: RocketCoordinateSystemOrientation = (
-        RocketCoordinateSystemOrientation.TAIL_TO_NOSE
-    )
+    coordinate_system_orientation: Literal['tail_to_nose', 'nose_to_tail'] = 'tail_to_nose'
     nose: NoseCone
     fins: List[Fins]
 

lib/models/sub/aerosurfaces.py

@@ -1,5 +1,4 @@
-from enum import Enum
-from typing import Optional, Tuple, List, Union
+from typing import Optional, Tuple, List, Union, Literal
 from pydantic import BaseModel, Field
 
 
@@ -28,18 +27,8 @@ class NoseCone(BaseModel):
     rocket_radius: float
 
 
-class FinsKinds(str, Enum):
-    TRAPEZOIDAL: str = "TRAPEZOIDAL"
-    ELLIPTICAL: str = "ELLIPTICAL"
-
-
-class AngleUnit(str, Enum):
-    RADIANS: str = "RADIANS"
-    DEGREES: str = "DEGREES"
-
-
 class Fins(BaseModel):
-    fins_kind: FinsKinds
+    fins_kind: Literal['trapezoidal', 'elliptical']
     name: str
     n: int
     root_chord: float
@@ -50,7 +39,7 @@ class Fins(BaseModel):
     tip_chord: Optional[float] = None
     cant_angle: Optional[float] = None
     rocket_radius: Optional[float] = None
-    airfoil: Optional[Tuple[List[Tuple[float, float]], AngleUnit]] = None
+    airfoil: Optional[Tuple[List[Tuple[float, float]], Literal['radians', 'degrees']]] = None
 
     def get_additional_parameters(self):
         return {

lib/models/sub/tanks.py

@@ -10,25 +10,11 @@ class TankKinds(str, Enum):
     ULLAGE: str = "ULLAGE"
 
 
-class TankCoordinateSystemOrientation(str, Enum):
-    NOZZLE_TO_COMBUSTION_CHAMBER: str = "NOZZLE_TO_COMBUSTION_CHAMBER"
-    COMBUSTION_CHAMBER_TO_NOZZLE: str = "COMBUSTION_CHAMBER_TO_NOZZLE"
-
-
 class TankFluids(BaseModel):
     name: str
     density: float
 
 
-class InterpolationMethods(str, Enum):
-    LINEAR: str = "LINEAR"
-    SPLINE: str = "SPLINE"
-    AKIMA: str = "AKIMA"
-    POLYNOMIAL: str = "POLYNOMIAL"
-    SHEPARD: str = "SHEPARD"
-    RBF: str = "RBF"
-
-
 class MotorTank(BaseModel):
     # Required parameters
     geometry: List[Tuple[Tuple[float, float], float]]

lib/routes/environment.py

@@ -116,13 +116,13 @@ async def read_rocketpy_env(environment_id: str):
 
 
 @router.get("/{environment_id}/summary")
-async def simulate_env(environment_id: str) -> EnvironmentSummary:
+async def get_environment_simulation(environment_id: str) -> EnvironmentSummary:
     """
     Loads rocketpy.environment simulation
 
     ## Args
     ``` environment_id: str ```
     """
-    with tracer.start_as_current_span("simulate_env"):
+    with tracer.start_as_current_span("get_environment_simulation"):
         controller = EnvironmentController()
         return await controller.get_environment_summary(environment_id)

lib/services/environment.py

@@ -29,7 +29,7 @@ def from_env_model(cls, env: EnvironmentModel) -> Self:
             date=env.date,
         )
         rocketpy_env.set_atmospheric_model(
-            type=env.atmospheric_model_type.value.lower(),
+            type=env.atmospheric_model_type,
             file=env.atmospheric_model_file,
         )
         return cls(environment=rocketpy_env)

lib/services/flight.py

@@ -34,7 +34,7 @@ def from_flight_model(cls, flight: FlightView) -> Self:
             rail_length=flight.rail_length,
             terminate_on_apogee=flight.terminate_on_apogee,
             time_overshoot=flight.time_overshoot,
-            equations_of_motion=flight.equations_of_motion.value.lower(),
+            equations_of_motion=flight.equations_of_motion,
             **flight.get_additional_parameters(),
         )
         return cls(flight=rocketpy_flight)

lib/services/motor.py

@@ -42,8 +42,8 @@ def from_motor_model(cls, motor: MotorView) -> Self:
             "dry_mass": motor.dry_mass,
             "dry_inertia": motor.dry_inertia,
             "center_of_dry_mass_position": motor.center_of_dry_mass_position,
-            "coordinate_system_orientation": motor.coordinate_system_orientation.value.lower(),
-            "interpolation_method": motor.interpolation_method.value.lower(),
+            "coordinate_system_orientation": motor.coordinate_system_orientation,
+            "interpolation_method": motor.interpolation_method,
             "reshape_thrust_curve": False or motor.reshape_thrust_curve,
         }
 

lib/services/rocket.py

@@ -43,7 +43,7 @@ def from_rocket_model(cls, rocket: RocketView) -> Self:
             power_off_drag=rocket.power_off_drag,
             power_on_drag=rocket.power_on_drag,
             center_of_mass_without_motor=rocket.center_of_mass_without_motor,
-            coordinate_system_orientation=rocket.coordinate_system_orientation.value.lower(),
+            coordinate_system_orientation=rocket.coordinate_system_orientation,
         )
 
         # RailButtons
@@ -157,15 +157,15 @@ def get_rocketpy_finset(fins: Fins, kind: str) -> RocketPyFins:
             RocketPyEllipticalFins
         """
         match kind:
-            case "TRAPEZOIDAL":
+            case "trapezoidal":
                 rocketpy_finset = RocketPyTrapezoidalFins(
                     n=fins.n,
                     name=fins.name,
                     root_chord=fins.root_chord,
                     span=fins.span,
                     **fins.get_additional_parameters(),
                 )
-            case "ELLIPTICAL":
+            case "elliptical":
                 rocketpy_finset = RocketPyEllipticalFins(
                     n=fins.n,
                     name=fins.name,

lib/views/motor.py

@@ -1,7 +1,6 @@
 from typing import List, Any, Optional
 from pydantic import BaseModel, ConfigDict
 from lib.views.interface import ApiBaseView
-from lib.models.sub.tanks import TankCoordinateSystemOrientation
 from lib.models.motor import MotorModel, MotorKinds
 from lib.utils import to_python_primitive
 
@@ -12,9 +11,7 @@ class MotorSummary(BaseModel):
     burn_out_time: Optional[float] = None
     burn_start_time: Optional[float] = None
     center_of_dry_mass_position: Optional[float] = None
-    coordinate_system_orientation: str = (
-        TankCoordinateSystemOrientation.NOZZLE_TO_COMBUSTION_CHAMBER.value
-    )
+    coordinate_system_orientation: str = 'nozzle_to_combustion_chamber'
     dry_I_11: Optional[float] = None
     dry_I_12: Optional[float] = None
     dry_I_13: Optional[float] = None

lib/views/rocket.py

@@ -1,16 +1,14 @@
 from typing import Any, Optional
 from pydantic import ConfigDict
-from lib.models.rocket import RocketModel, RocketCoordinateSystemOrientation
+from lib.models.rocket import RocketModel
 from lib.views.interface import ApiBaseView
 from lib.views.motor import MotorView, MotorSummary
 from lib.utils import to_python_primitive
 
 
 class RocketSummary(MotorSummary):
     area: Optional[float] = None
-    coordinate_system_orientation: str = (
-        RocketCoordinateSystemOrientation.TAIL_TO_NOSE.value
-    )
+    coordinate_system_orientation: str = 'tail_to_nose'
     center_of_mass_without_motor: Optional[float] = None
     motor_center_of_dry_mass_position: Optional[float] = None
     motor_position: Optional[float] = None

tests/test_routes/conftest.py

@@ -9,14 +9,14 @@
 
 
 @pytest.fixture
-def stub_env():
+def stub_environment_dump():
     env = EnvironmentModel(latitude=0, longitude=0)
     env_json = env.model_dump_json()
     return json.loads(env_json)
 
 
 @pytest.fixture
-def stub_motor():
+def stub_motor_dump():
     motor = MotorModel(
         thrust_source=[[0, 0]],
         burn_time=0,
@@ -30,7 +30,7 @@ def stub_motor():
 
 
 @pytest.fixture
-def stub_tank():
+def stub_tank_dump():
     tank = MotorTank(
         geometry=[[(0, 0), 0]],
         gas=TankFluids(name='gas', density=0),
@@ -81,7 +81,7 @@ def stub_mass_tank(stub_tank):
 
 
 @pytest.fixture
-def stub_nose_cone():
+def stub_nose_cone_dump():
     nose_cone = NoseCone(
         name='nose',
         length=0,
@@ -95,9 +95,9 @@ def stub_nose_cone():
 
 
 @pytest.fixture
-def stub_fins():
+def stub_fins_dump():
     fins = Fins(
-        fins_kind='TRAPEZOIDAL',
+        fins_kind='trapezoidal',
         name='fins',
         n=0,
         root_chord=0,
@@ -109,7 +109,7 @@ def stub_fins():
 
 
 @pytest.fixture
-def stub_rocket(stub_motor, stub_nose_cone, stub_fins):
+def stub_rocket_dump(stub_motor, stub_nose_cone, stub_fins):
     rocket = RocketModel(
         motor=stub_motor,
         radius=0,
@@ -121,7 +121,7 @@ def stub_rocket(stub_motor, stub_nose_cone, stub_fins):
         power_on_drag=[(0, 0)],
         nose=stub_nose_cone,
         fins=[stub_fins],
-        coordinate_system_orientation='TAIL_TO_NOSE',
+        coordinate_system_orientation='tail_to_nose',
     )
     rocket_json = rocket.model_dump_json()
     return json.loads(rocket_json)