Extend sensor data context type

tum-esm · Jan 7, 2024 · f46fec1 · f46fec1
1 parent 814fa09
commit f46fec1
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 7 deletions.
diff --git a/em27_metadata/interfaces.py b/em27_metadata/interfaces.py
@@ -135,6 +135,18 @@ def get(
                         self.locations.root,
                     )
                 )
+                atmospheric_profile_location: em27_metadata.types.LocationMetadata
+                if setup.atmospheric_profile_location_id is not None:
+                    atmospheric_profile_location = next(
+                        filter(
+                            lambda l: l.location_id == setup.
+                            atmospheric_profile_location_id,
+                            self.locations.root,
+                        )
+                    )
+                else:
+                    atmospheric_profile_location = location
+
             except StopIteration:
                 continue
             try:
@@ -162,11 +174,7 @@ def get(
                         if setup.pressure_data_source else sensor.sensor_id
                     ),
                     calibration_factors=calibration_factor,
-                    atmospheric_profile_location_id=(
-                        setup.atmospheric_profile_location_id
-                        if setup.atmospheric_profile_location_id else
-                        setup.location_id
-                    ),
+                    atmospheric_profile_location=atmospheric_profile_location,
                 )
             )
 

diff --git a/em27_metadata/types.py b/em27_metadata/types.py
@@ -257,7 +257,7 @@ class SensorDataContext(pydantic.BaseModel):
     utc_offset: float
     pressure_data_source: str
     calibration_factors: CalibrationFactors
-    atmospheric_profile_location_id: str
+    atmospheric_profile_location: LocationMetadata
 
     @pydantic.field_serializer("from_datetime", "to_datetime")
     def t_serializer(self, dt: datetime.date, _info: Any) -> str:

diff --git a/tests/test_getter_function.py b/tests/test_getter_function.py
@@ -147,7 +147,7 @@ def test_getter_function() -> None:
     assert location_ids == ["lid1", "lid1", "lid2", "lid2", "lid2", "lid2"]
 
     atmospheric_profile_location_ids = [
-        c.atmospheric_profile_location_id for c in chunks
+        c.atmospheric_profile_location.location_id for c in chunks
     ]
     assert atmospheric_profile_location_ids == [
         "lid1", "lid1", "lid2", "lid2", "lid2", "lid1"