@@ -54,10 +54,8 @@ def for_flight(cls) -> 'DiscretizeConfig':
5454
5555
5656class InfinityEncoder (RocketPyEncoder ):
57- def __init__ (self , * args , ** kwargs ):
58- super ().__init__ (* args , ** kwargs )
59-
60- def default (self , obj ):
57+ def default (self , o ):
58+ obj = o
6159 if (
6260 isinstance (obj , Function )
6361 and not callable (obj .source )
@@ -75,7 +73,7 @@ def default(self, obj):
7573 mutate_self = False ,
7674 )
7775 if isinstance (obj , Flight ):
78- obj ._Flight__evaluate_post_process
76+ obj ._Flight__evaluate_post_process ()
7977 solution = np .array (obj .solution )
8078 size = len (solution )
8179 if size > 25 :
@@ -117,92 +115,80 @@ def rocketpy_encoder(obj):
117115
118116
119117def collect_attributes (obj , attribute_classes = None ):
120- """
121- Collect attributes from various simulation classes and populate them from the flight object.
122- """
123- if attribute_classes is None :
124- attribute_classes = []
118+ """Collect and serialize attributes from simulation classes."""
119+ attribute_classes = attribute_classes or []
125120
126121 attributes = rocketpy_encoder (obj )
127-
128122 for attribute_class in attribute_classes :
129- if issubclass (attribute_class , FlightSimulation ):
130- flight_attributes_list = [
131- attr
132- for attr in attribute_class .__annotations__ .keys ()
133- if attr not in "message" , "rocket" , "env" ]
134- ]
135- try :
136- for key in flight_attributes_list :
137- if key not in attributes :
138- try :
139- value = getattr (obj , key )
140- attributes [key ] = value
141- except Exception :
142- pass
143- except Exception :
144- pass
145-
146- elif issubclass (attribute_class , RocketSimulation ):
147- rocket_attributes_list = [
148- attr
149- for attr in attribute_class .__annotations__ .keys ()
150- if attr not in "message" , "motor" ]
151- ]
152- try :
153- for key in rocket_attributes_list :
154- if key not in attributes .get ("rocket" , {}):
155- try :
156- value = getattr (obj .rocket , key )
157- attributes .setdefault ("rocket" , {})[key ] = value
158- except Exception :
159- pass
160- except Exception :
161- pass
162-
163- elif issubclass (attribute_class , MotorSimulation ):
164- motor_attributes_list = [
165- attr
166- for attr in attribute_class .__annotations__ .keys ()
167- if attr not in "message" ]
168- ]
169- try :
170- for key in motor_attributes_list :
171- if key not in attributes .get ("rocket" , {}).get (
172- "motor" , {}
173- ):
174- try :
175- value = getattr (obj .rocket .motor , key )
176- attributes .setdefault ("rocket" , {}).setdefault (
177- "motor" , {}
178- )[key ] = value
179- except Exception :
180- pass
181- except Exception :
182- pass
183-
184- elif issubclass (attribute_class , EnvironmentSimulation ):
185- environment_attributes_list = [
186- attr
187- for attr in attribute_class .__annotations__ .keys ()
188- if attr not in "message" ]
189- ]
190- try :
191- for key in environment_attributes_list :
192- if key not in attributes .get ("env" , {}):
193- try :
194- value = getattr (obj .env , key )
195- attributes .setdefault ("env" , {})[key ] = value
196- except Exception :
197- pass
198- except Exception :
199- pass
200- else :
201- continue
123+ _populate_simulation_attributes (obj , attribute_class , attributes )
202124
203125 return rocketpy_encoder (attributes )
204126
205127
128+ def _populate_simulation_attributes (obj , attribute_class , attributes ):
129+ if not isinstance (attribute_class , type ):
130+ return
131+
132+ mappings = (
133+ (FlightSimulation , (), (), {"message" , "rocket" , "env" }),
134+ (RocketSimulation , ("rocket" ,), ("rocket" ,), {"message" , "motor" }),
135+ (
136+ MotorSimulation ,
137+ ("rocket" , "motor" ),
138+ ("rocket" , "motor" ),
139+ {"message" },
140+ ),
141+ (EnvironmentSimulation , ("env" ,), ("env" ,), {"message" }),
142+ )
143+
144+ for klass , source_path , target_path , exclusions in mappings :
145+ if not issubclass (attribute_class , klass ):
146+ continue
147+
148+ keys = _annotation_keys (attribute_class , exclusions )
149+ if not keys :
150+ return
151+
152+ source = _resolve_attribute_path (obj , source_path )
153+ if source is None :
154+ return
155+
156+ target = _resolve_attribute_target (attributes , target_path )
157+ _copy_missing_attributes (source , target , keys )
158+ return
159+
160+
161+ def _annotation_keys (attribute_class , exclusions ):
162+ annotations = getattr (attribute_class , "__annotations__" , {})
163+ return [key for key in annotations if key not in exclusions ]
164+
165+
166+ def _resolve_attribute_path (root , path ):
167+ current = root
168+ for attr in path :
169+ if current is None :
170+ return None
171+ current = getattr (current , attr , None )
172+ return current
173+
174+
175+ def _resolve_attribute_target (attributes , path ):
176+ target = attributes
177+ for key in path :
178+ target = target .setdefault (key , {})
179+ return target
180+
181+
182+ def _copy_missing_attributes (source , target , keys ):
183+ for key in keys :
184+ if key in target :
185+ continue
186+ try :
187+ target [key ] = getattr (source , key )
188+ except AttributeError :
189+ continue
190+
191+
206192def _fix_datetime_fields (data ):
207193 """
208194 Fix datetime fields that RocketPyEncoder converted to lists.
0 commit comments