diff --git a/icoscp_stilt/src/icoscp_stilt/stilt.py b/icoscp_stilt/src/icoscp_stilt/stilt.py index 812e878..a45d041 100644 --- a/icoscp_stilt/src/icoscp_stilt/stilt.py +++ b/icoscp_stilt/src/icoscp_stilt/stilt.py @@ -74,6 +74,15 @@ def has_observation_data(self) -> bool: """ return self.icosId is not None +@dataclass(frozen=True) +class ObservationResult: + dobj: DataObjectLite + columns: ArraysDict + +@dataclass(frozen=True) +class ObservationDfResult: + dobj: DataObjectLite + df: pd.DataFrame def list_stations() -> list[StiltStation]: """ @@ -244,7 +253,7 @@ def fetch_observations_pandas( spec: URL, stations: list[StiltStation], columns: list[str] | None = None -) -> dict[str, pd.DataFrame]: +) -> dict[str, ObservationDfResult]: """ Batch-fetch observational datasets for a number of STILT stations @@ -257,19 +266,19 @@ def fetch_observations_pandas( within the observational datasets; if omitted, all columns are included in the result - :return (dict[str, DataFrame]): a dictionary with STILT station IDs as - keys and pandas DataFrames with observational datasets as values + :return (dict[str, ObservationDfResult]): a dictionary with STILT station IDs as + keys and instances of ObservationDfResult as values """ return { - st_id: pd.DataFrame(arrs) - for st_id, arrs in fetch_observations(spec, stations, columns).items() + st_id: ObservationDfResult(obs_res.dobj, pd.DataFrame(obs_res.columns)) + for st_id, obs_res in fetch_observations(spec, stations, columns).items() } def fetch_observations( spec: URL, stations: list[StiltStation], columns: list[str] | None = None -) -> dict[str, ArraysDict]: +) -> dict[str, ObservationResult]: """ Batch-fetch observational datasets for a number of STILT stations. Is a lower-level and better-performing version of `fetch_observations_pandas`. @@ -287,8 +296,8 @@ def fetch_observations( within the observational datasets; if omitted, all columns are included in the result - :return (dict[str, dict[str, ndarray]]): a dictionary with STILT station IDs as - keys and as values, dictionaries of column names vs numpy arrays + :return (dict[str, ObservationResult]): a dictionary with STILT station IDs as + keys and ObservationResult instances as values """ icos2ss: dict[tuple[URL, float], StiltStation] = { @@ -312,8 +321,8 @@ def lookup_ss(dobj: DataObjectLite) -> list[StiltStation]: if lookup_ss(dobj) ] return { - ss.id: arrs - for dobj, arrs in data.batch_get_columns_as_arrays(dobjs_to_fetch, columns) + ss.id: ObservationResult(dobj, cols) + for dobj, cols in data.batch_get_columns_as_arrays(dobjs_to_fetch, columns) for ss in lookup_ss(dobj) }