From 0f352597e59bf14df4a311cd40eb7dbc59249b0b Mon Sep 17 00:00:00 2001
From: peterdudfield <peter.dudfield@hotmail.com>
Date: Thu, 31 Aug 2023 10:53:49 +0100
Subject: [PATCH] lint

---
 src/database.py        | 2 +-
 src/pydantic_models.py | 9 ++++++++-
 src/utils.py           | 9 ++++++++-
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/database.py b/src/database.py
index 2e35ef53..6152b1da 100644
--- a/src/database.py
+++ b/src/database.py
@@ -40,7 +40,7 @@
     convert_forecasts_to_many_datetime_many_generation,
     convert_location_sql_to_many_datetime_many_generation,
 )
-from utils import floor_30_minutes_dt, get_start_datetime, filter_forecast_values
+from utils import floor_30_minutes_dt, filter_forecast_values, get_start_datetime
 
 logger = structlog.stdlib.get_logger()
 
diff --git a/src/pydantic_models.py b/src/pydantic_models.py
index cbbc6d11..e74e3fa1 100644
--- a/src/pydantic_models.py
+++ b/src/pydantic_models.py
@@ -109,7 +109,10 @@ def convert_location_sql_to_many_datetime_many_generation(
 
 
 def convert_forecasts_to_many_datetime_many_generation(
-    forecasts: List[ForecastSQL], historic: bool = True
+    forecasts: List[ForecastSQL],
+    historic: bool = True,
+    start_datetime_utc: Optional[datetime] = None,
+    end_datetime_utc: Optional[datetime] = None,
 ) -> List[OneDatetimeManyForecastValues]:
     """Change forecasts to list of OneDatetimeManyForecastValues
 
@@ -134,6 +137,10 @@ def convert_forecasts_to_many_datetime_many_generation(
 
         for forecast_value in forecast_values:
             datetime_utc = forecast_value.target_time
+            if start_datetime_utc is not None and datetime_utc < start_datetime_utc:
+                continue
+            if end_datetime_utc is not None and datetime_utc > end_datetime_utc:
+                continue
             forecast_mw = round(forecast_value.expected_power_generation_megawatts, 2)
 
             # if the datetime object is not in the dictionary, add it
diff --git a/src/utils.py b/src/utils.py
index 41a31d64..a4f42297 100644
--- a/src/utils.py
+++ b/src/utils.py
@@ -5,7 +5,7 @@
 
 import numpy as np
 import structlog
-from nowcasting_datamodel.models import ForecastValue, ForecastSQL
+from nowcasting_datamodel.models import ForecastSQL, ForecastValue
 from pydantic import Field
 from pytz import timezone
 
@@ -67,6 +67,13 @@ def floor_6_hours_dt(dt: datetime):
 
 
 def format_datetime(datetime_str: str = None):
+    """
+    Format datetime string to datetime object
+
+    If None return None, if not timezone, add UTC
+    :param datetime_str:
+    :return:
+    """
     if datetime_str is None:
         return None