diff --git a/tests/test_sdk.py b/tests/test_sdk.py index 9ce4b77..dcb87f7 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -208,6 +208,17 @@ def test_get_historical_pandas(self): self.assertIn("point_time", df.columns) self.assertIn("value", df.columns) + def test_get_historical_pandas_include_imputed(self): + start = datetime.now() - timedelta(days=7) + end = datetime.now() + df = self.historical.get_historical_pandas(start, end, REGION, include_imputed_marker=True) + + self.assertIsInstance(df, pd.DataFrame) + self.assertGreaterEqual(len(df), 1) + self.assertIn("point_time", df.columns) + self.assertIn("value", df.columns) + self.assertIn("imputed_data_used", df.columns) + def test_get_historical_pandas_meta(self): start = datetime.now() - timedelta(days=7) end = datetime.now() diff --git a/watttime/api.py b/watttime/api.py index 199e1cc..fadcad9 100644 --- a/watttime/api.py +++ b/watttime/api.py @@ -370,6 +370,7 @@ def get_historical_jsons( Literal["co2_moer", "co2_aoer", "health_damage"] ] = "co2_moer", model: Optional[Union[str, date]] = None, + include_imputed_marker: bool = False, ) -> List[dict]: """ Base function to scrape historical data, returning a list of .json responses. @@ -391,6 +392,9 @@ def get_historical_jsons( url = "{}/v3/historical".format(self.url_base) params = {"region": region, "signal_type": signal_type} + if include_imputed_marker: + params["include_imputed_marker"] = "true" + start, end = self._parse_dates(start, end) chunks = self._get_chunks(start, end) @@ -421,6 +425,7 @@ def get_historical_pandas( ] = "co2_moer", model: Optional[Union[str, date]] = None, include_meta: bool = False, + include_imputed_marker: bool = False, ): """ Return a pd.DataFrame with point_time, and values. @@ -434,7 +439,9 @@ def get_historical_pandas( Returns: pd.DataFrame: _description_ """ - responses = self.get_historical_jsons(start, end, region, signal_type, model) + responses = self.get_historical_jsons( + start, end, region, signal_type, model, include_imputed_marker + ) df = pd.json_normalize( responses, record_path="data", meta=["meta"] if include_meta else [] ) @@ -452,6 +459,7 @@ def get_historical_csv( Literal["co2_moer", "co2_aoer", "health_damage"] ] = "co2_moer", model: Optional[Union[str, date]] = None, + include_imputed_marker: bool = False, ): """ Retrieves historical data from a specified start date to an end date and saves it as a CSV file. @@ -467,7 +475,9 @@ def get_historical_csv( Returns: None, results are saved to a csv file in the user's home directory. """ - df = self.get_historical_pandas(start, end, region, signal_type, model) + df = self.get_historical_pandas( + start, end, region, signal_type, model, include_imputed_marker + ) out_dir = Path.home() / "watttime_historical_csvs" out_dir.mkdir(exist_ok=True)