ColocatedDataLists class

metno · Aug 7, 2023 · 7e2da85 · 7e2da85
1 parent e6837ab
commit 7e2da85
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/pyaerocom/colocation_3d.py b/pyaerocom/colocation_3d.py
@@ -46,9 +46,14 @@
 logger = logging.getLogger(__name__)
 
 
-ColocatedDataLists = namedtuple(
-    "ColocatedDataLists", ["colocateddata_for_statistics", "colocateddata_for_profile_viz"]
-)
+# ColocatedDataLists = namedtuple(
+#    "ColocatedDataLists", ["colocateddata_for_statistics", "colocateddata_for_profile_viz"]
+# )
+
+
+class ColocatedDataLists(NamedTuple):
+    colocateddata_for_statistics: list[ColocatedData]
+    colocateddata_for_profile_viz: list[ColocatedData]
 
 
 def _colocate_vertical_profile_gridded(
@@ -332,7 +337,7 @@ def colocate_vertical_profile_gridded(
     colocation_layer_limits: list[dict] = None,
     profile_layer_limits: list[dict] = None,
     **kwargs,
-) -> NamedTuple:
+) -> ColocatedDataLists:
     """
     Colocated vertical profile data with gridded (model) data
 
@@ -492,7 +497,7 @@ def colocate_vertical_profile_gridded(
         coldata.data.attrs["just_for_viz"] = 1
 
     colocated_data_lists = ColocatedDataLists(
-        output_prep[0], output_prep[1]
+        *output_prep
     )  # put the list of prepared output into namedtuple object s.t. both position and named arguments can be used
     # breakpoint()
     return colocated_data_lists
diff --git a/pyaerocom/colocation_auto.py b/pyaerocom/colocation_auto.py
@@ -1477,7 +1477,7 @@ def _run_helper(self, model_var: str, obs_var: str):
             if self.save_coldata:
                 self._save_coldata(coldata)
 
-        elif isinstance(coldata, ColocatedDataLists):
+        elif isinstance(coldata, ColocatedDataLists):  # look into intertools chain.from_iterable
             for i_list in coldata:
                 for coldata_obj in i_list:
                     coldata_obj.data.attrs["model_name"] = self.get_model_name()