add timezone info if start or end date are None

onnela-lab · Oct 17, 2023 · 1661710 · 1661710
1 parent 958fcce
commit 1661710
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/forest/oak/base.py b/forest/oak/base.py
@@ -497,13 +497,17 @@ def run(study_folder: str, output_folder: str, tz_str: str = None,
             date_start = date_start - timedelta(hours=date_start.hour)
         else:
             date_start = datetime.strptime(time_start, fmt)
+            date_start = date_start.replace(
+                tzinfo=from_zone
+            ).astimezone(to_zone)
             date_start = date_start - timedelta(hours=date_start.hour)
 
         if time_end is None:
             date_end = dates_shifted[-1]
             date_end = date_end - timedelta(hours=date_end.hour)
         else:
             date_end = datetime.strptime(time_end, fmt)
+            date_end = date_end.replace(tzinfo=from_zone).astimezone(to_zone)
             date_end = date_end - timedelta(hours=date_end.hour)
 
         days = pd.date_range(date_start, date_end, freq='D')