ENH: refactor attribute collection to use unified collect_attributes function

Author: aasitvora99

src/services/environment.py

@@ -5,7 +5,7 @@
 from rocketpy.environment.environment import Environment as RocketPyEnvironment
 from src.models.environment import EnvironmentModel
 from src.views.environment import EnvironmentSimulation
-from src.utils import collect_simulation_attributes
+from src.utils import collect_attributes
 
 
 class EnvironmentService:
@@ -54,9 +54,9 @@ def get_environment_simulation(self) -> EnvironmentSimulation:
             EnvironmentSimulation
         """
 
-        encoded_attributes = collect_simulation_attributes(
+        encoded_attributes = collect_attributes(
             self.environment,
-            EnvironmentSimulation,
+            [EnvironmentSimulation],
         )
         env_simulation = EnvironmentSimulation(**encoded_attributes)
         return env_simulation

src/services/flight.py

@@ -11,7 +11,7 @@
 from src.views.rocket import RocketSimulation
 from src.views.motor import MotorSimulation
 from src.views.environment import EnvironmentSimulation
-from src.utils import collect_simulation_attributes
+from src.utils import collect_attributes
 
 class FlightService:
     _flight: RocketPyFlight
@@ -57,12 +57,9 @@ def get_flight_simulation(self) -> FlightSimulation:
         Returns:
             FlightSimulation
         """
-        encoded_attributes = collect_simulation_attributes(
+        encoded_attributes = collect_attributes(
             self.flight,
-            FlightSimulation,
-            RocketSimulation, 
-            MotorSimulation,
-            EnvironmentSimulation
+            [FlightSimulation, RocketSimulation, MotorSimulation, EnvironmentSimulation]
         )
         flight_simulation = FlightSimulation(**encoded_attributes)
         return flight_simulation

src/services/motor.py

@@ -17,7 +17,7 @@
 from src.models.sub.tanks import TankKinds
 from src.models.motor import MotorKinds, MotorModel
 from src.views.motor import MotorSimulation
-from src.utils import collect_simulation_attributes
+from src.utils import collect_attributes
 
 
 class MotorService:
@@ -140,9 +140,9 @@ def get_motor_simulation(self) -> MotorSimulation:
         Returns:
             MotorSimulation
         """
-        encoded_attributes = collect_simulation_attributes(
+        encoded_attributes = collect_attributes(
             self.motor,
-            MotorSimulation,
+            [MotorSimulation],
         )
         motor_simulation = MotorSimulation(**encoded_attributes)
         return motor_simulation

src/services/rocket.py

@@ -18,7 +18,7 @@
 from src.services.motor import MotorService
 from src.views.rocket import RocketSimulation
 from src.views.motor import MotorSimulation
-from src.utils import collect_simulation_attributes
+from src.utils import collect_attributes
 
 class RocketService:
     _rocket: RocketPyRocket
@@ -107,10 +107,9 @@ def get_rocket_simulation(self) -> RocketSimulation:
         Returns:
             RocketSimulation
         """
-        encoded_attributes = collect_simulation_attributes(
+        encoded_attributes = collect_attributes(
             self.rocket,
-            RocketSimulation, 
-            MotorSimulation,
+            [RocketSimulation, MotorSimulation]
         )
         rocket_simulation = RocketSimulation(**encoded_attributes)
         return rocket_simulation

src/utils.py

@@ -5,6 +5,11 @@
 
 from typing import NoReturn
 
+from views.environment import EnvironmentSimulation
+from views.flight import FlightSimulation
+from views.motor import MotorSimulation
+from views.rocket import RocketSimulation
+
 from rocketpy._encoders import RocketPyEncoder
 from starlette.datastructures import Headers, MutableHeaders
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
@@ -37,93 +42,107 @@ def rocketpy_encoder(obj):
     return json.loads(json_str)
 
 
-def collect_simulation_attributes(flight_obj, flight_simulation_class, rocket_simulation_class, motor_simulation_class, environment_simulation_class):
+def collect_attributes(obj, attribute_classes=None):
     """
     Collect attributes from various simulation classes and populate them from the flight object.
     
     Args:
-        flight_obj: RocketPy Flight object
-        flight_simulation_class: FlightSimulation class
-        rocket_simulation_class: RocketSimulation class  
-        motor_simulation_class: MotorSimulation class
-        environment_simulation_class: EnvironmentSimulation class
-    
+        obj: RocketPy Flight object
+        attribute_classes: List of attribute classes to collect from
+
     Returns:
         Dictionary with all collected attributes
     """
-    attributes = rocketpy_encoder(flight_obj)
-    
-    flight_attributes_list = [
-        attr for attr in flight_simulation_class.__annotations__.keys() 
-        if attr not in ['message', 'rocket', 'env']
-    ]
-
-    rocket_attributes_list = [
-        attr for attr in rocket_simulation_class.__annotations__.keys() 
-        if attr not in ['message', 'motor']
-    ]
-
-    motor_attributes_list = [
-        attr for attr in motor_simulation_class.__annotations__.keys() 
-        if attr not in ['message']
-    ]
-
-    environment_attributes_list = [
-        attr for attr in environment_simulation_class.__annotations__.keys()
-        if attr not in ['message']
-    ]
-
-    try:
-        for key in flight_attributes_list:
+    if attribute_classes is None:
+        attribute_classes = []
+
+    attributes = rocketpy_encoder(obj)
+
+    for attribute_class in attribute_classes:
+        if issubclass(attribute_class, FlightSimulation):
+            flight_attributes_list = [
+                attr for attr in attribute_class.__annotations__.keys() 
+                if attr not in ['message', 'rocket', 'env']
+            ]
             try:
-                value = getattr(flight_obj, key)
-                attributes[key] = value
-            except AttributeError as e:
-                logger.warning(f"Flight attribute '{key}' not found: {e}")
-            except Exception as e:
-                logger.error(f"Error getting flight attribute '{key}': {type(e).__name__}: {e}")
-    except Exception as e:
-        logger.error(f"Error processing flight attributes: {type(e).__name__}: {e}")
-
-    try:
-        for key in rocket_attributes_list:
+                for key in flight_attributes_list:
+                    if key not in attributes:
+                        try:
+                            value = getattr(obj, key)
+                            attributes[key] = value
+                        except AttributeError:
+                            pass
+                        except Exception:
+                            pass
+            except Exception:
+                pass
+                
+        elif issubclass(attribute_class, RocketSimulation):
+            rocket_attributes_list = [
+                attr for attr in attribute_class.__annotations__.keys() 
+                if attr not in ['message', 'motor']
+            ]
             try:
-                value = getattr(flight_obj.rocket, key)
-                attributes["rocket"][key] = value
-            except AttributeError as e:
-                logger.warning(f"Rocket attribute '{key}' not found: {e}")
-            except Exception as e:
-                logger.error(f"Error getting rocket attribute '{key}': {type(e).__name__}: {e}")
-    except Exception as e:
-        logger.error(f"Error processing rocket attributes: {type(e).__name__}: {e}")
-
-    try:
-        for key in motor_attributes_list:
+                for key in rocket_attributes_list:
+                    if key not in attributes.get("rocket", {}):
+                        try:
+                            value = getattr(obj.rocket, key)
+                            if "rocket" not in attributes:
+                                attributes["rocket"] = {}
+                            attributes["rocket"][key] = value
+                        except AttributeError:
+                            pass
+                        except Exception:
+                            pass
+            except Exception:
+                pass
+                
+        elif issubclass(attribute_class, MotorSimulation):
+            motor_attributes_list = [
+                attr for attr in attribute_class.__annotations__.keys() 
+                if attr not in ['message']
+            ]
             try:
-                value = getattr(flight_obj.rocket.motor, key)
-                attributes["rocket"]["motor"][key] = value
-            except AttributeError as e:
-                logger.warning(f"Motor attribute '{key}' not found: {e}")
-            except Exception as e:
-                logger.error(f"Error getting motor attribute '{key}': {type(e).__name__}: {e}")
-    except Exception as e:
-        logger.error(f"Error processing motor attributes: {type(e).__name__}: {e}")
-    
-    try:
-        for key in environment_attributes_list:
+                for key in motor_attributes_list:
+                    if key not in attributes.get("rocket", {}).get("motor", {}):
+                        try:
+                            value = getattr(obj.rocket.motor, key)
+                            if "rocket" not in attributes:
+                                attributes["rocket"] = {}
+                            if "motor" not in attributes["rocket"]:
+                                attributes["rocket"]["motor"] = {}
+                            attributes["rocket"]["motor"][key] = value
+                        except AttributeError:
+                            pass
+                        except Exception:
+                            pass
+            except Exception:
+                pass
+                
+        elif issubclass(attribute_class, EnvironmentSimulation):
+            environment_attributes_list = [
+                attr for attr in attribute_class.__annotations__.keys()
+                if attr not in ['message']
+            ]
             try:
-                value = getattr(flight_obj.env, key)
-                attributes["env"][key] = value
-            except AttributeError as e:
-                logger.warning(f"Environment attribute '{key}' not found: {e}")
-            except Exception as e:
-                logger.error(f"Error getting environment attribute '{key}': {type(e).__name__}: {e}")
-    except Exception as e:
-        logger.error(f"Error processing environment attributes: {type(e).__name__}: {e}")
+                for key in environment_attributes_list:
+                    if key not in attributes.get("env", {}):
+                        try:
+                            value = getattr(obj.env, key)
+                            if "env" not in attributes:
+                                attributes["env"] = {}
+                            attributes["env"][key] = value
+                        except AttributeError:
+                            pass
+                        except Exception:
+                            pass
+            except Exception:
+                pass
+        else:
+            continue
 
     return rocketpy_encoder(attributes)
 
-
 class RocketPyGZipMiddleware:
     def __init__(
         self, app: ASGIApp, minimum_size: int = 500, compresslevel: int = 9