From e658c3036fb9e75ac34164aa422c3a22c36b0955 Mon Sep 17 00:00:00 2001
From: vitor-aignosi <vitor.santos@aignosi.com.br>
Date: Thu, 16 Apr 2026 09:20:04 -0300
Subject: [PATCH] feat: add timezone handling for DataFrame index

- Implemented _set_timezone_on_index method to ensure DataFrame indices are set to UTC if not already timezone-aware.
- Updated training and validation DataFrame processing to include timezone configuration for improved data consistency.
---
 .../repository/data_manager_repository.py     | 32 ++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/model_manager/utils/repository/data_manager_repository.py b/model_manager/utils/repository/data_manager_repository.py
index 88e35b8..0406761 100644
--- a/model_manager/utils/repository/data_manager_repository.py
+++ b/model_manager/utils/repository/data_manager_repository.py
@@ -158,7 +158,8 @@ class DataManagerRepository(SientiaMonitoring):
 
         train_df = _ensure_date_column_parsed(train_df, params)
         train_df = self._configure_datetime_index(train_df, params, metadata)
-
+        train_df = self._set_timezone_on_index(train_df, metadata)
+        
         if len(train_df) <= 0:
             raise ValueError('Training data view is empty after transformation')
 
@@ -179,6 +180,7 @@ class DataManagerRepository(SientiaMonitoring):
 
             val_df = _ensure_date_column_parsed(val_df, params)
             val_df = self._configure_datetime_index(val_df, params, metadata)
+            val_df = self._set_timezone_on_index(val_df, metadata)
 
             if len(val_df) <= 0:
                 raise ValueError('Validation data view is empty after transformation')
@@ -350,6 +352,13 @@ class DataManagerRepository(SientiaMonitoring):
 
         Guards against None to avoid 'NoneType' object has no attribute 'index' downstream.
         Prefers params.date_column when set; otherwise looks for common timestamp column names.
+        Args:
+            data: The DataFrame to configure the datetime index for.
+            params: The training parameters.
+            metadata: The metadata for the training run.
+
+        Returns:
+            The DataFrame with the datetime index configured.
         """
         if data is None:
             raise ValueError(
@@ -409,6 +418,27 @@ class DataManagerRepository(SientiaMonitoring):
         )
         return data
 
+    def _set_timezone_on_index(self, data: pd.DataFrame, metadata: dict[str, Any] | None = None) -> pd.DataFrame:
+        """
+        Check if the index has a timezone and if not, set it to UTC timezone.
+
+        Args:
+            data: The DataFrame to set the timezone on.
+            metadata: The metadata for the training run.
+
+        Returns:
+            The DataFrame with the timezone set.
+        """
+
+        if isinstance(data.index, pd.DatetimeIndex):
+            if data.index.tz is None:
+                data.index = data.index.tz_localize('UTC')
+            else:
+                data.index = data.index.tz_convert('UTC')
+        else:
+            raise ValueError('Index is not a DatetimeIndex')
+                
+        return data
     def _get_reports_directory(self) -> str:
         """
         Get the absolute path to the reports directory.