More handling of null end timestamps

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