Only allow a single model to be returned through /historical (#14)

tests/test_sdk.py

@@ -29,7 +29,7 @@ def __init__(self, json_data, status_code):
 
         def json(self):
             return self.json_data
-        
+
         def raise_for_status(self):
             assert self.status_code == 200
 
@@ -190,6 +190,8 @@ def test_get_historical_pandas_meta(self):
         self.assertIn("value", df.columns)
         self.assertIn("meta", df.columns)
 
+        assert pd.api.types.is_datetime64_any_dtype(df["point_time"].dtype)
+
     def test_get_historical_csv(self):
         start = parse("2022-01-01 00:00Z")
         end = parse("2022-01-02 00:00Z")
@@ -203,6 +205,34 @@ def test_get_historical_csv(self):
         assert fp.exists()
         fp.unlink()
 
+    def test_multi_model_range(self):
+        """If model is not specified, we should only return the most recent model data"""
+        myaccess = WattTimeMyAccess()
+        access = myaccess.get_access_pandas()
+        access = access.loc[
+            (access["signal_type"] == "co2_moer") & (access["region"] == REGION)
+        ].sort_values("model", ascending=False)
+        assert len(access) > 1
+
+        # start request one month before data_start of most recent model
+        start = access["data_start"].values[0] - pd.Timedelta(days=30)
+        end = access["data_start"].values[0] + pd.Timedelta(days=30)
+        df = self.historical.get_historical_pandas(
+            start, end, REGION, include_meta=True
+        )
+
+        # should not span into an older model
+        self.assertEqual(df.iloc[0]["meta"]["model"]["date"], access.iloc[0]["model"])
+
+        self.assertEqual(df.iloc[-1]["meta"]["model"]["date"], access.iloc[0]["model"])
+
+        # first point_time should be data_start from my-acces
+        self.assertAlmostEqual(
+            df.iloc[0]["point_time"],
+            access.iloc[0]["data_start"].tz_localize("UTC"),
+            delta=pd.Timedelta(days=1),
+        )
+
 
 class TestWattTimeMyAccess(unittest.TestCase):
     def setUp(self):
@@ -257,6 +287,10 @@ def test_access_pandas(self):
         self.assertIn("type", df.columns)
         self.assertGreaterEqual(len(df), 1)
 
+        assert pd.api.types.is_datetime64_any_dtype(df["data_start"])
+        assert pd.api.types.is_datetime64_any_dtype(df["train_start"])
+        assert pd.api.types.is_datetime64_any_dtype(df["train_end"])
+
 
 class TestWattTimeForecast(unittest.TestCase):
     def setUp(self):
@@ -301,28 +335,27 @@ def test_historical_forecast_pandas(self):
         self.assertIn("point_time", df.columns)
         self.assertIn("value", df.columns)
         self.assertIn("generated_at", df.columns)
-        
+
     def test_horizon_hours(self):
         json = self.forecast.get_forecast_json(region=REGION, horizon_hours=0)
         self.assertIsInstance(json, dict)
         self.assertIn("meta", json)
         self.assertEqual(len(json["data"]), 1)
         self.assertIn("point_time", json["data"][0])
-        
+
         json2 = self.forecast.get_forecast_json(region=REGION, horizon_hours=24)
         self.assertIsInstance(json2, dict)
         self.assertIn("meta", json2)
         self.assertEqual(len(json2["data"]), 288)
         self.assertIn("point_time", json2["data"][0])
-                
+
         json3 = self.forecast.get_forecast_json(region=REGION, horizon_hours=72)
         self.assertIsInstance(json3, dict)
         self.assertIn("meta", json3)
         self.assertEqual(len(json3["data"]), 864)
         self.assertIn("point_time", json3["data"][0])
 
 
-
 class TestWattTimeMaps(unittest.TestCase):
     def setUp(self):
         self.maps = WattTimeMaps()
@@ -353,9 +386,11 @@ def test_get_maps_json_health(self):
             parse(health["meta"]["last_updated"]), parse("2022-01-01 00:00Z")
         )
         self.assertGreater(len(health["features"]), 100)  # 114 as of 2023-12-01
-        
+
     def test_region_from_loc(self):
-        region = self.maps.region_from_loc(latitude=39.7522, longitude=-105.0, signal_type='co2_moer')
+        region = self.maps.region_from_loc(
+            latitude=39.7522, longitude=-105.0, signal_type="co2_moer"
+        )
         self.assertEqual(region["region"], "PSCO")
         self.assertEqual(region["region_full_name"], "Public Service Co of Colorado")
         self.assertEqual(region["signal_type"], "co2_moer")

watttime/api.py

@@ -223,6 +223,14 @@ def get_historical_jsons(
             if len(j["meta"]["warnings"]):
                 print("\n", "Warnings Returned:", params, j["meta"])
 
+        # the API should not let this happen, but ensure for sanity
+        unique_models = set([r["meta"]["model"]["date"] for r in responses])
+        chosen_model = model_date or max(unique_models)
+        if len(unique_models) > 1:
+            responses = [
+                r for r in responses if r["meta"]["model"]["date"] == chosen_model
+            ]
+
         return responses
 
     def get_historical_pandas(
@@ -254,6 +262,9 @@ def get_historical_pandas(
         df = pd.json_normalize(
             responses, record_path="data", meta=["meta"] if include_meta else []
         )
+
+        df["point_time"] = pd.to_datetime(df["point_time"])
+
         return df
 
     def get_historical_csv(
@@ -340,7 +351,14 @@ def get_access_pandas(self) -> pd.DataFrame:
                             }
                         )
 
-        return pd.DataFrame(out)
+        out = pd.DataFrame(out)
+        out = out.assign(
+            data_start=pd.to_datetime(out["data_start"]),
+            train_start=pd.to_datetime(out["train_start"]),
+            train_end=pd.to_datetime(out["train_end"]),
+        )
+
+        return out
 
 
 class WattTimeForecast(WattTimeBase):
@@ -451,7 +469,7 @@ def get_historical_forecast_json(
         params = {
             "region": region,
             "signal_type": signal_type,
-            horizon_hours: horizon_hours,
+            "horizon_hours": horizon_hours,
         }
 
         start, end = self._parse_dates(start, end)
@@ -495,7 +513,7 @@ def get_historical_forecast_pandas(
             start (Union[str, datetime]): The start date or datetime for the historical forecast.
             end (Union[str, datetime]): The end date or datetime for the historical forecast.
             region (str): The region for which the historical forecast data is retrieved.
-            signal_type (Optional[Literal["co2_moer", "co2_aoer", "health_damage"]], optional): 
+            signal_type (Optional[Literal["co2_moer", "co2_aoer", "health_damage"]], optional):
                 The type of signal for the historical forecast data. Defaults to "co2_moer".
             model_date (Optional[Union[str, date]], optional): The model date for the historical forecast data. Defaults to None.
             horizon_hours (int, optional): The number of hours to forecast. Defaults to 24. Minimum of 0 provides a "nowcast" created with the forecast, maximum of 72.