Feat: Add helper function for replace timezone (#428)

* feat: add small helper function for replacing timezone

* refactor: use helper function for timezone replacement

* refactor: use helper function in fip mapper as well

* test: update tests

* fix linting

* more linting fixes

* refactor: use more focused function

* test: update tests

* fix linting

Author: arjunsridhar12345

src/aind_metadata_mapper/fip/mapper.py

@@ -57,6 +57,7 @@
     get_intended_measurements,
     get_procedures,
     get_protocols_for_modality,
+    normalize_utc_timezone,
 )
 
 logger = logging.getLogger(__name__)
@@ -302,8 +303,10 @@ def transform(
 
         # Get timing from all data streams (handle multiple epochs)
         # Find earliest start_time and latest end_time across all epochs
-        start_times = [ensure_timezone(ds["start_time"]) for ds in data_streams]
-        end_times = [ensure_timezone(ds["end_time"]) for ds in data_streams]
+        start_times = [
+            ensure_timezone(normalize_utc_timezone(ds["start_time"])) for ds in data_streams
+        ]
+        end_times = [ensure_timezone(normalize_utc_timezone(ds["end_time"])) for ds in data_streams]
 
         earliest_start = min(start_times)
         latest_end = max(end_times)

src/aind_metadata_mapper/gather_metadata.py

@@ -25,7 +25,7 @@
 from aind_metadata_mapper.base import MapperJobSettings
 from aind_metadata_mapper.mapper_registry import registry
 from aind_metadata_mapper.models import JobSettings
-from aind_metadata_mapper.utils import get_procedures, get_subject, metadata_service_helper
+from aind_metadata_mapper.utils import get_procedures, get_subject, metadata_service_helper, normalize_utc_timezone
 
 
 class GatherMetadataJob:
@@ -184,7 +184,9 @@ def build_data_description(self, acquisition_start_time: str, subject_id: str) -
             logging.debug(f"Using existing {file_name}.")
             return self._get_file_from_user_defined_directory(file_name=file_name)
 
-        acquisition_start_time.replace("Z", "+00:00")  # remove when we're past Python 3.11
+        acquisition_start_time = normalize_utc_timezone(
+            acquisition_start_time
+        )  # remove when we're past Python 3.11
         creation_time = datetime.fromisoformat(acquisition_start_time)
         logging.info(f"Using acquisition start time: {creation_time}")
 
@@ -511,7 +513,9 @@ def _validate_acquisition_start_time(self, acquisition_start_time: str) -> str:
         ValueError
             If acquisition_start_time doesn't match settings and raise_if_invalid is True
         """
-        acquisition_start_time = acquisition_start_time.replace("Z", "+00:00")  # remove when we're past Python 3.11
+        acquisition_start_time = normalize_utc_timezone(
+            acquisition_start_time
+        )  # remove when we're past Python 3.11
         local_acq_start_time = datetime.fromisoformat(acquisition_start_time)
 
         if self.settings.acquisition_start_time and local_acq_start_time != self.settings.acquisition_start_time:

src/aind_metadata_mapper/utils.py

@@ -20,6 +20,36 @@
 SUBJECT_BASE_URL = "http://aind-metadata-service/api/v2/subject"
 
 
+def normalize_utc_timezone(dt: str) -> str:
+    """
+    Normalize UTC timezone indicator from 'Z' to '+00:00' offset format.
+
+    This ensures compatibility with Python 3.10's datetime.fromisoformat(),
+    which doesn't support the 'Z' shorthand for UTC.
+
+    Parameters
+    ----------
+    dt : str
+        An ISO 8601 datetime string, potentially ending with 'Z'.
+
+    Returns
+    -------
+    str
+        The datetime string with 'Z' replaced by '+00:00' if present,
+        otherwise unchanged.
+
+    Examples
+    --------
+    >>> normalize_utc_timezone("2025-11-16T23:00:22Z")
+    '2025-11-16T23:00:22+00:00'
+    >>> normalize_utc_timezone("2025-11-16T23:00:22-05:00")
+    '2025-11-16T23:00:22-05:00'
+    """
+    if dt.endswith("Z"):
+        return dt[:-1] + "+00:00"
+    return dt
+
+
 def ensure_timezone(dt):
     """Ensure datetime has timezone info using system local timezone.
 
@@ -126,7 +156,8 @@ def get_procedures(subject_id: str, base_url: str = PROCEDURES_BASE_URL) -> Opti
 
 
 def get_intended_measurements(
-    subject_id: str, base_url: str = "http://aind-metadata-service/intended_measurements"
+    subject_id: str,
+    base_url: str = "http://aind-metadata-service/intended_measurements",
 ) -> Optional[dict]:
     """Fetch intended measurements for a subject from the metadata service.
 
@@ -245,7 +276,10 @@ def get_instrument(
                 )
             return None
         else:
-            return sorted(matching_records, key=lambda record: record["modification_date"])[-1]
+            return sorted(
+                matching_records,
+                key=lambda record: record["modification_date"],
+            )[-1]
     except Exception as e:
         logger.warning(f"Unexpected error fetching instrument {instrument_id}: {e}")
         return None
@@ -319,7 +353,11 @@ def check_existing_instrument(
     bool
         True if a record with the same instrument_id and modification_date exists.
     """
-    existing = get_instrument(instrument_id, modification_date=modification_date, suppress_warning=True)
+    existing = get_instrument(
+        instrument_id,
+        modification_date=modification_date,
+        suppress_warning=True,
+    )
     return existing is not None
 
 

tests/test_utils.py

@@ -23,12 +23,21 @@
     get_protocols_for_modality,
     get_subject,
     prompt_for_string,
+    normalize_utc_timezone,
 )
 
 
 class TestUtils(unittest.TestCase):
     """Test cases for utility functions in aind_metadata_mapper."""
 
+    def test_replace_z_with_offset(self):
+        """Test that a trailing 'Z' timezone shorthand is replaced with '+00:00'."""
+        self.assertEqual(normalize_utc_timezone("2025-11-16T23:00:22Z"), "2025-11-16T23:00:22+00:00")
+
+    def test_no_replacement(self):
+        """Test that the original string is returned unchanged when 'old' is not found."""
+        self.assertEqual(normalize_utc_timezone("2025-11-16T23:00:22"), "2025-11-16T23:00:22")
+
     def test_ensure_timezone_none(self):
         """Test that ensure_timezone handles None input by returning current time with timezone.