Skip to content

Commit bcca778

Browse files
GabrielBarberiniGabrielBarberini
committed
migrate Enum to Literal when possible
1 parent 4876a5c commit bcca778

File tree

18 files changed

+505
-684
lines changed

18 files changed

+505
-684
lines changed

‎lib/controllers/environment.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ async def get_rocketpy_env_binary(
4141
return env_service.get_env_binary()
4242

4343
@controller_exception_handler
44-
async def simulate_env(
44+
async def get_environment_simulation(
4545
self, env_id: str
4646
) -> EnvironmentSummary:
4747
"""

‎lib/models/environment.py‎

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,8 @@
11
import datetime
2-
from enum import Enum
3-
from typing import Optional, ClassVar, Self
2+
from typing import Optional, ClassVar, Self, Literal
43
from lib.models.interface import ApiBaseModel
54

65

7-
class AtmosphericModelTypes(str, Enum):
8-
STANDARD_ATMOSPHERE: str = "STANDARD_ATMOSPHERE"
9-
CUSTOM_ATMOSPHERE: str = "CUSTOM_ATMOSPHERE"
10-
WYOMING_SOUNDING: str = "WYOMING_SOUNDING"
11-
FORECAST: str = "FORECAST"
12-
REANALYSIS: str = "REANALYSIS"
13-
ENSEMBLE: str = "ENSEMBLE"
14-
15-
166
class EnvironmentModel(ApiBaseModel):
177
NAME: ClassVar = 'environment'
188
METHODS: ClassVar = ('POST', 'GET', 'PUT', 'DELETE')
@@ -21,9 +11,7 @@ class EnvironmentModel(ApiBaseModel):
2111
elevation: Optional[int] = 1
2212

2313
# Optional parameters
24-
atmospheric_model_type: AtmosphericModelTypes = (
25-
AtmosphericModelTypes.STANDARD_ATMOSPHERE
26-
)
14+
atmospheric_model_type: Literal['standard_atmosphere', 'custom_atmosphere', 'wyoming_sounding', 'forecast', 'reanalysis', 'ensemble'] = 'standard_atmosphere'
2715
atmospheric_model_file: Optional[str] = None
2816
date: Optional[datetime.datetime] = (
2917
datetime.datetime.today() + datetime.timedelta(days=1)

‎lib/models/flight.py‎

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,9 @@
1-
from enum import Enum
2-
from typing import Optional, Self, ClassVar
1+
from typing import Optional, Self, ClassVar, Literal
32
from lib.models.interface import ApiBaseModel
43
from lib.models.rocket import RocketModel
54
from lib.models.environment import EnvironmentModel
65

76

8-
class EquationsOfMotion(str, Enum):
9-
STANDARD: str = "STANDARD"
10-
SOLID_PROPULSION: str = "SOLID_PROPULSION"
11-
12-
137
class FlightModel(ApiBaseModel):
148
NAME: ClassVar = "flight"
159
METHODS: ClassVar = ("POST", "GET", "PUT", "DELETE")
@@ -20,7 +14,7 @@ class FlightModel(ApiBaseModel):
2014
rail_length: float = 1
2115
time_overshoot: bool = True
2216
terminate_on_apogee: bool = True
23-
equations_of_motion: EquationsOfMotion = EquationsOfMotion.STANDARD
17+
equations_of_motion: Literal['standard', 'solid_propulsion'] = 'standard'
2418

2519
# Optional parameters
2620
inclination: Optional[int] = None

‎lib/models/motor.py‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from pydantic import PrivateAttr, model_validator
44

55
from lib.models.interface import ApiBaseModel
6-
from lib.models.sub.tanks import MotorTank, InterpolationMethods, TankCoordinateSystemOrientation
6+
from lib.models.sub.tanks import MotorTank
77

88

99
class MotorKinds(str, Enum):
@@ -48,10 +48,8 @@ class MotorModel(ApiBaseModel):
4848
throat_radius: Optional[float] = None
4949

5050
# Optional parameters
51-
interpolation_method: InterpolationMethods = InterpolationMethods.LINEAR
52-
coordinate_system_orientation: TankCoordinateSystemOrientation = (
53-
TankCoordinateSystemOrientation.NOZZLE_TO_COMBUSTION_CHAMBER
54-
)
51+
interpolation_method: Literal['linear', 'spline', 'akima', 'polynomial', 'shepard', 'rbf'] = 'linear'
52+
coordinate_system_orientation: Literal['nozzle_to_combustion_chamber', 'combustion_chamber_to_nozzle'] = 'nose_to_combustion_chamber'
5553
reshape_thrust_curve: Union[bool, tuple] = False
5654

5755
# Computed parameters

‎lib/models/rocket.py‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
1-
from enum import Enum
2-
from typing import Optional, Tuple, List, Union, Self, ClassVar
1+
from typing import Optional, Tuple, List, Union, Self, ClassVar, Literal
32
from lib.models.interface import ApiBaseModel
43
from lib.models.motor import MotorModel
54
from lib.models.sub.aerosurfaces import (
@@ -11,11 +10,6 @@
1110
)
1211

1312

14-
class RocketCoordinateSystemOrientation(str, Enum):
15-
TAIL_TO_NOSE: str = "TAIL_TO_NOSE"
16-
NOSE_TO_TAIL: str = "NOSE_TO_TAIL"
17-
18-
1913
class RocketModel(ApiBaseModel):
2014
NAME: ClassVar = "rocket"
2115
METHODS: ClassVar = ("POST", "GET", "PUT", "DELETE")
@@ -32,9 +26,7 @@ class RocketModel(ApiBaseModel):
3226
] = (0, 0, 0)
3327
power_off_drag: List[Tuple[float, float]] = [(0, 0)]
3428
power_on_drag: List[Tuple[float, float]] = [(0, 0)]
35-
coordinate_system_orientation: RocketCoordinateSystemOrientation = (
36-
RocketCoordinateSystemOrientation.TAIL_TO_NOSE
37-
)
29+
coordinate_system_orientation: Literal['tail_to_nose', 'nose_to_tail'] = 'tail_to_nose'
3830
nose: NoseCone
3931
fins: List[Fins]
4032

‎lib/models/sub/aerosurfaces.py‎

Lines changed: 3 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
1-
from enum import Enum
2-
from typing import Optional, Tuple, List, Union
1+
from typing import Optional, Tuple, List, Union, Literal
32
from pydantic import BaseModel, Field
43

54

@@ -28,18 +27,8 @@ class NoseCone(BaseModel):
2827
rocket_radius: float
2928

3029

31-
class FinsKinds(str, Enum):
32-
TRAPEZOIDAL: str = "TRAPEZOIDAL"
33-
ELLIPTICAL: str = "ELLIPTICAL"
34-
35-
36-
class AngleUnit(str, Enum):
37-
RADIANS: str = "RADIANS"
38-
DEGREES: str = "DEGREES"
39-
40-
4130
class Fins(BaseModel):
42-
fins_kind: FinsKinds
31+
fins_kind: Literal['trapezoidal', 'elliptical']
4332
name: str
4433
n: int
4534
root_chord: float
@@ -50,7 +39,7 @@ class Fins(BaseModel):
5039
tip_chord: Optional[float] = None
5140
cant_angle: Optional[float] = None
5241
rocket_radius: Optional[float] = None
53-
airfoil: Optional[Tuple[List[Tuple[float, float]], AngleUnit]] = None
42+
airfoil: Optional[Tuple[List[Tuple[float, float]], Literal['radians', 'degrees']]] = None
5443

5544
def get_additional_parameters(self):
5645
return {

‎lib/models/sub/tanks.py‎

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -10,25 +10,11 @@ class TankKinds(str, Enum):
1010
ULLAGE: str = "ULLAGE"
1111

1212

13-
class TankCoordinateSystemOrientation(str, Enum):
14-
NOZZLE_TO_COMBUSTION_CHAMBER: str = "NOZZLE_TO_COMBUSTION_CHAMBER"
15-
COMBUSTION_CHAMBER_TO_NOZZLE: str = "COMBUSTION_CHAMBER_TO_NOZZLE"
16-
17-
1813
class TankFluids(BaseModel):
1914
name: str
2015
density: float
2116

2217

23-
class InterpolationMethods(str, Enum):
24-
LINEAR: str = "LINEAR"
25-
SPLINE: str = "SPLINE"
26-
AKIMA: str = "AKIMA"
27-
POLYNOMIAL: str = "POLYNOMIAL"
28-
SHEPARD: str = "SHEPARD"
29-
RBF: str = "RBF"
30-
31-
3218
class MotorTank(BaseModel):
3319
# Required parameters
3420
geometry: List[Tuple[Tuple[float, float], float]]

‎lib/routes/environment.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,13 +116,13 @@ async def read_rocketpy_env(environment_id: str):
116116

117117

118118
@router.get("/{environment_id}/summary")
119-
async def simulate_env(environment_id: str) -> EnvironmentSummary:
119+
async def get_environment_simulation(environment_id: str) -> EnvironmentSummary:
120120
"""
121121
Loads rocketpy.environment simulation
122122
123123
## Args
124124
``` environment_id: str ```
125125
"""
126-
with tracer.start_as_current_span("simulate_env"):
126+
with tracer.start_as_current_span("get_environment_simulation"):
127127
controller = EnvironmentController()
128128
return await controller.get_environment_summary(environment_id)

‎lib/services/environment.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ def from_env_model(cls, env: EnvironmentModel) -> Self:
2929
date=env.date,
3030
)
3131
rocketpy_env.set_atmospheric_model(
32-
type=env.atmospheric_model_type.value.lower(),
32+
type=env.atmospheric_model_type,
3333
file=env.atmospheric_model_file,
3434
)
3535
return cls(environment=rocketpy_env)

‎lib/services/flight.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ def from_flight_model(cls, flight: FlightView) -> Self:
3434
rail_length=flight.rail_length,
3535
terminate_on_apogee=flight.terminate_on_apogee,
3636
time_overshoot=flight.time_overshoot,
37-
equations_of_motion=flight.equations_of_motion.value.lower(),
37+
equations_of_motion=flight.equations_of_motion,
3838
**flight.get_additional_parameters(),
3939
)
4040
return cls(flight=rocketpy_flight)

0 commit comments

Comments
 (0)