diff --git a/src/cct_connector/ServiceAlertFixer.py b/src/cct_connector/ServiceAlertFixer.py index 15e66b9..0077c73 100644 --- a/src/cct_connector/ServiceAlertFixer.py +++ b/src/cct_connector/ServiceAlertFixer.py @@ -5,6 +5,7 @@ from db_utils import minio_utils import pandas as pd import pyarrow.dataset as ds +import pytz from cct_connector import ServiceAlertBase from cct_connector import ( @@ -16,6 +17,7 @@ SN_REGEX_PATTERN = r"^\d{10}$" SN_RE = re.compile(SN_REGEX_PATTERN) HM_RE = re.compile(r'^\d{2}:\d{2}$') +SAST_TZ = pytz.timezone('Africa/Johannesburg') def _clean_sa_df(data_df: pd.DataFrame) -> pd.DataFrame: @@ -45,7 +47,7 @@ def _clean_sa_df(data_df: pd.DataFrame) -> pd.DataFrame: lambda val: ( datetime.strptime( val.replace("60", "59") + "+02:00", "%H:%M%z" - ) if pd.notna(None) and HM_RE.match(val) else None + ) if (pd.notna(val) and HM_RE.match(val)) else None ) ), # Creating timestamps @@ -54,11 +56,11 @@ def _clean_sa_df(data_df: pd.DataFrame) -> pd.DataFrame: axis=1 ).dt.tz_localize("+02:00"), "forecast_end_timestamp": lambda df: df.apply( - lambda row: datetime.combine( + lambda row: SAST_TZ.localize(datetime.combine( # Assuming that it ends on the day of expiry (row["expiry_date"] - pd.Timedelta(days=1)).date(), - row["forecast_end_time"].dt.time - ).tz_localize("+02:00") if pd.notna(row['forecast_end_time']) else None, + row["forecast_end_time"].time() + )) if pd.notna(row['forecast_end_time']) else None, axis=1 ), "location": lambda df: df.apply( @@ -70,8 +72,7 @@ def _clean_sa_df(data_df: pd.DataFrame) -> pd.DataFrame: row["Address_x0020_Location_x0020_2"][:len(row["Description12"])] != row["Description12"][:len(row["Address_x0020_Location_x0020_2"])] ) else None - ), - axis=1 + ), axis=1 ) }).assign(**{ # fixing cases where the start and end timestamps roll over the day