Add optional pint unit conversion and per-client fixture tests

metloom/pointdata/base.py

@@ -1,12 +1,13 @@
 import copy
 import logging
 from datetime import datetime
-from typing import List
+from typing import Dict, List, Union
 
 import pandas as pd
 
 import geopandas as gpd
 
+from ..unit_conversions import convert_series
 from ..variables import SensorDescription, VariableBase
 
 LOG = logging.getLogger("metloom.pointdata.base")
@@ -168,6 +169,62 @@ def validate_sensor_df(cls, gdf: gpd.GeoDataFrame):
             if "_units" not in rc:
                 assert f"{rc}_units" in remaining_columns
 
+    def _convert_units(
+        self,
+        gdf: gpd.GeoDataFrame,
+        desired_units: Union[str, Dict[str, str], None],
+    ):
+        """
+        Optionally convert variable columns to a desired unit using pint.
+
+        This is a no-op unless ``desired_units`` is provided - when it is
+        ``None`` pint is never invoked and the dataframe is returned unchanged.
+
+        Args:
+            gdf: the assembled GeoDataFrame returned by a ``get_*`` method.
+                May be ``None`` or empty.
+            desired_units: either a single unit string applied to every
+                variable column, or a dict mapping ``variable.name`` to a unit
+                string. Units strings are pint-compatible (e.g. ``"degC"``,
+                ``"mm"``, ``"W/m^2"``).
+        Returns:
+            The GeoDataFrame with the requested variable columns converted in
+            place. Converted values are plain magnitudes (never pint
+            Quantities) and the matching ``{name}_units`` column is updated to
+            the target unit. NaNs are preserved.
+        """
+        # Short-circuit before touching pint when no conversion was requested
+        if desired_units is None or gdf is None or len(gdf) == 0:
+            return gdf
+
+        # variable columns are those that have a sibling `{col}_units` column
+        variable_columns = [
+            c for c in gdf.columns if f"{c}_units" in gdf.columns
+        ]
+        for column in variable_columns:
+            if isinstance(desired_units, str):
+                target = desired_units
+            else:
+                target = desired_units.get(column)
+            if target is None:
+                continue
+
+            units_column = f"{column}_units"
+            # infer the current unit from the first non-null entry
+            current = gdf[units_column].dropna()
+            if len(current) == 0:
+                LOG.warning(
+                    "No units found for %s; skipping conversion", column
+                )
+                continue
+            source_unit = current.iloc[0]
+
+            gdf[column] = convert_series(
+                gdf[column].to_numpy(), source_unit, target
+            )
+            gdf[units_column] = target
+        return gdf
+
     def __repr__(self):
         return f"{self.__class__.__name__}({self.id!r}, {self.name!r})"
 
@@ -197,6 +254,7 @@ def get_daily_data(
         start_date: datetime,
         end_date: datetime,
         variables: List[SensorDescription],
+        desired_units: Union[str, Dict[str, str]] = None,
     ):
         """
         Get daily measurement data
@@ -205,6 +263,10 @@ def get_daily_data(
             end_date: datetime object for end of data collection period
             variables: List of metloom.variables.SensorDescription object
                 from self.ALLOWED_VARIABLES
+            desired_units: Optional pint-compatible unit conversion. Either a
+                single unit string applied to every variable, or a dict
+                mapping ``variable.name`` to a unit string. When omitted no
+                conversion is performed and the inferred units are returned.
         Returns:
             GeoDataFrame of data. The dataframe should be indexed on
             ['datetime', 'site'] and have columns
@@ -222,6 +284,7 @@ def get_hourly_data(
         start_date: datetime,
         end_date: datetime,
         variables: List[SensorDescription],
+        desired_units: Union[str, Dict[str, str]] = None,
     ):
         """
         Get hourly measurement data
@@ -230,6 +293,10 @@ def get_hourly_data(
             end_date: datetime object for end of data collection period
             variables: List of metloom.variables.SensorDescription object
                 from self.ALLOWED_VARIABLES
+            desired_units: Optional pint-compatible unit conversion. Either a
+                single unit string applied to every variable, or a dict
+                mapping ``variable.name`` to a unit string. When omitted no
+                conversion is performed and the inferred units are returned.
         Returns:
             GeoDataFrame of data. The dataframe should be indexed on
             ['datetime', 'site'] and have columns
@@ -247,6 +314,7 @@ def get_snow_course_data(
         start_date: datetime,
         end_date: datetime,
         variables: List[SensorDescription],
+        desired_units: Union[str, Dict[str, str]] = None,
     ):
         """
         Get snow course data
@@ -255,6 +323,10 @@ def get_snow_course_data(
             end_date: datetime object for end of data collection period
             variables: List of metloom.variables.SensorDescription object
                 from self.ALLOWED_VARIABLES
+            desired_units: Optional pint-compatible unit conversion. Either a
+                single unit string applied to every variable, or a dict
+                mapping ``variable.name`` to a unit string. When omitted no
+                conversion is performed and the inferred units are returned.
         Returns:
             GeoDataFrame of data. The dataframe should be indexed on
             ['datetime', 'site'] and have columns

metloom/pointdata/cdec.py

@@ -228,7 +228,8 @@ def _get_data(
         end_date: datetime,
         variables: List[SensorDescription],
         duration_list: List[str],
-        include_measurement_date=False
+        include_measurement_date=False,
+        desired_units=None,
     ):
         """
         Args:
@@ -240,6 +241,8 @@ def _get_data(
             include_measurement_date: boolean for including the
                 'measurmentDate' column in the resulting dataframe. This column
                 is only relevant for snow courses
+            desired_units: Optional pint-compatible unit conversion (see
+                PointData.get_daily_data)
         Returns:
             GeoDataFrame of data, indexed on datetime, site
         """
@@ -280,54 +283,68 @@ def _get_data(
             else:
                 df = None
         self.validate_sensor_df(df)
+        df = self._convert_units(df, desired_units)
         return df
 
     def get_event_data(
         self,
         start_date: datetime,
         end_date: datetime,
         variables: List[SensorDescription],
+        desired_units=None,
     ):
-        return self._get_data(start_date, end_date, variables, ["E"])
+        return self._get_data(
+            start_date, end_date, variables, ["E"], desired_units=desired_units
+        )
 
     def get_daily_data(
         self,
         start_date: datetime,
         end_date: datetime,
         variables: List[SensorDescription],
+        desired_units=None,
     ):
         """
         See docstring for PointData.get_daily_data
         Example query:
         https://cdec.water.ca.gov/dynamicapp/req/JSONDataServlet?
         Stations=TNY&SensorNums=3&dur_code=D&Start=2021-05-16&End=2021-05-16
         """
-        return self._get_data(start_date, end_date, variables, ["D", "h", "E"])
+        return self._get_data(
+            start_date, end_date, variables, ["D", "h", "E"],
+            desired_units=desired_units,
+        )
 
     def get_hourly_data(
         self,
         start_date: datetime,
         end_date: datetime,
         variables: List[SensorDescription],
+        desired_units=None,
     ):
         """
         See docstring for PointData.get_hourly_data
         """
-        return self._get_data(start_date, end_date, variables, ["h", "E"])
+        return self._get_data(
+            start_date, end_date, variables, ["h", "E"],
+            desired_units=desired_units,
+        )
 
     def get_snow_course_data(
         self,
         start_date: datetime,
         end_date: datetime,
         variables: List[SensorDescription],
+        desired_units=None,
     ):
         """
         See docstring for PointData.get_snow_course_data
         """
         if not self.is_partly_snow_course():
             raise ValueError(f"{self.id} is not a snow course")
         return self._get_data(
-            start_date, end_date, variables, ["M"], include_measurement_date=True
+            start_date, end_date, variables, ["M"],
+            include_measurement_date=True, desired_units=desired_units,
         )
 
     @staticmethod

metloom/pointdata/cues.py

@@ -114,7 +114,7 @@ def _sensor_response_to_df(self, data, variable):
 
     def _get_data(
         self, start_date, end_date, variables: List[SensorDescription],
-        period,
+        period, desired_units=None,
     ):
         df = pd.DataFrame()
         df.index.name = "datetime"
@@ -133,18 +133,23 @@ def _get_data(
         df = gpd.GeoDataFrame(df, geometry=[self.metadata] * len(df))
         df = df.reset_index().set_index(["datetime", "site"])
         self.validate_sensor_df(df)
+        df = self._convert_units(df, desired_units)
         return df
 
     def get_daily_data(self, start_date: datetime, end_date: datetime,
-                       variables: List[SensorDescription]):
+                       variables: List[SensorDescription],
+                       desired_units=None):
         return self._get_data(
-            start_date, end_date, variables, "day"
+            start_date, end_date, variables, "day",
+            desired_units=desired_units,
         )
 
     def get_hourly_data(self, start_date: datetime, end_date: datetime,
-                        variables: List[SensorDescription]):
+                        variables: List[SensorDescription],
+                        desired_units=None):
         return self._get_data(
-            start_date, end_date, variables, "hr"
+            start_date, end_date, variables, "hr",
+            desired_units=desired_units,
         )
 
     def get_snow_course_data(self, start_date: datetime, end_date: datetime,

metloom/pointdata/files.py

@@ -169,7 +169,7 @@ def _get_one_variable(self, resp_df, period, variable: SensorDescription):
         return data
 
     def _get_data(self, start_date, end_date, variables: List[SensorDescription],
-                  period):
+                  period, desired_units=None):
         """
         Utilizes cached data or downloads the data
         """
@@ -222,18 +222,23 @@ def _get_data(self, start_date, end_date, variables: List[SensorDescription],
         df = df.reset_index().set_index(["datetime", "site"])
 
         self.validate_sensor_df(df)
+        df = self._convert_units(df, desired_units)
         return df
 
     def get_daily_data(self, start_date: datetime, end_date: datetime,
-                       variables: List[SensorDescription]):
+                       variables: List[SensorDescription],
+                       desired_units=None):
         return self._get_data(
-            start_date, end_date, variables, "D"
+            start_date, end_date, variables, "D",
+            desired_units=desired_units,
         )
 
     def get_hourly_data(self, start_date: datetime, end_date: datetime,
-                        variables: List[SensorDescription]):
+                        variables: List[SensorDescription],
+                        desired_units=None):
         return self._get_data(
-            start_date, end_date, variables, "h"
+            start_date, end_date, variables, "h",
+            desired_units=desired_units,
         )
 
     def _get_metadata(self):

metloom/pointdata/geosphere_austria.py

@@ -172,6 +172,7 @@ def _get_data(
         end_date: datetime,
         variables: List[SensorDescription],
         desired_duration: str,
+        desired_units=None,
     ):
         """
 
@@ -181,6 +182,8 @@ def _get_data(
             variables: List of metloom.variables.SensorDescription object
                 from self.ALLOWED_VARIABLES
             desired_duration: duration code ['D', 'h']
+            desired_units: Optional pint-compatible unit conversion (see
+                PointData.get_daily_data)
         Returns:
             GeoDataFrame of data, indexed on datetime, site
         """
@@ -212,6 +215,7 @@ def _get_data(
             else:
                 df = None
         self.validate_sensor_df(df)
+        df = self._convert_units(df, desired_units)
         return df
 
     @classmethod
@@ -346,6 +350,7 @@ def get_daily_data(
         start_date: datetime,
         end_date: datetime,
         variables: List[SensorDescription],
+        desired_units=None,
     ):
         """
         See docstring for PointData.get_daily_data
@@ -354,19 +359,24 @@ def get_daily_data(
         tawes-v1-10min?parameters=TL&station_ids=11035
         """
         self._validate_dates(end_date)
-        return self._get_data(start_date, end_date, variables, "D")
+        return self._get_data(
+            start_date, end_date, variables, "D", desired_units=desired_units
+        )
 
     def get_hourly_data(
         self,
         start_date: datetime,
         end_date: datetime,
         variables: List[SensorDescription],
+        desired_units=None,
     ):
         """
         See docstring for PointData.get_hourly_data
         """
         self._validate_dates(end_date)
-        return self._get_data(start_date, end_date, variables, "h")
+        return self._get_data(
+            start_date, end_date, variables, "h", desired_units=desired_units
+        )
 
 
 class GeoSphereHistPointData(GeoSpherePointDataBase):
@@ -406,11 +416,14 @@ def get_daily_data(
         start_date: datetime,
         end_date: datetime,
         variables: List[SensorDescription],
+        desired_units=None,
     ):
         """
         See docstring for PointData.get_daily_data
         Example query:
         https://dataset.api.hub.geosphere.at/v1/station/historical/klima-v1-1d
         ?station_ids=11401&start=2023-04-12&end=2023-04-14&parameters=schnee
         """
-        return self._get_data(start_date, end_date, variables, None)
+        return self._get_data(
+            start_date, end_date, variables, None, desired_units=desired_units
+        )