Accept the Z designator in timestamptz strings

pyiceberg/utils/datetime.py

@@ -31,8 +31,9 @@
 ISO_TIMESTAMP = re.compile(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(.\d{1,6})?")
 ISO_TIMESTAMP_NANO = re.compile(r"(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(.\d{1,6})?(\d{1,3})?")
 EPOCH_TIMESTAMPTZ = datetime.fromisoformat("1970-01-01T00:00:00.000000+00:00")
-ISO_TIMESTAMPTZ = re.compile(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(.\d{1,6})?[-+]\d{2}:\d{2}")
-ISO_TIMESTAMPTZ_NANO = re.compile(r"(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(.\d{1,6})?(\d{1,3})?([-+]\d{2}:\d{2})")
+# Seconds are optional and the zone can be an offset or the UTC designator, as in Java's ISO_DATE_TIME
+ISO_TIMESTAMPTZ = re.compile(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}(?::\d{2}(.\d{1,6})?)?([-+]\d{2}:\d{2}|[Zz])")
+ISO_TIMESTAMPTZ_NANO = re.compile(r"(\d{4}-\d{2}-\d{2}T\d{2}:\d{2})(?:(:\d{2})(.\d{1,6})?(\d{1,3})?)?([-+]\d{2}:\d{2}|[Zz])")
 
 
 def micros_to_days(timestamp: int) -> int:
@@ -133,16 +134,22 @@ def timestamp_to_nanos(timestamp_str: str) -> int:
     raise ValueError(f"Invalid timestamp without zone: {timestamp_str} (must be ISO-8601)")
 
 
+def _utc_offset(offset: str) -> str:
+    # datetime.fromisoformat only accepts the "Z" designator from Python 3.11
+    return "+00:00" if offset in ("Z", "z") else offset
+
+
 def timestamptz_to_nanos(timestamptz_str: str) -> int:
     """Convert an ISO-8601 formatted timestamp with zone to nanoseconds from 1970-01-01T00:00:00.000000000+00:00."""
     if match := ISO_TIMESTAMPTZ_NANO.fullmatch(timestamptz_str):
         # Python datetime does not have native nanoseconds support
         # Hence we need to extract nanoseconds timestamp manually
-        # group(3) holds the sub-microsecond digits (fraction positions 7-9), so
+        # group(4) holds the sub-microsecond digits (fraction positions 7-9), so
         # right-pad to 3 digits before reading them as nanoseconds (e.g. "7" -> 700).
-        ns_str = (match.group(3) or "0").ljust(3, "0")
-        ms_str = match.group(2) if match.group(2) else ""
-        timestamptz_str_without_ns_str = match.group(1) + ms_str + match.group(4)
+        ns_str = (match.group(4) or "0").ljust(3, "0")
+        seconds_str = match.group(2) or ""
+        ms_str = match.group(3) or ""
+        timestamptz_str_without_ns_str = match.group(1) + seconds_str + ms_str + _utc_offset(match.group(5))
         return datetime_to_nanos(datetime.fromisoformat(timestamptz_str_without_ns_str)) + int(ns_str)
     if ISO_TIMESTAMP_NANO.fullmatch(timestamptz_str):
         # When we can match a timestamp without a zone, we can give a more specific error
@@ -167,7 +174,9 @@ def millis_to_datetime(millis: int) -> datetime:
 
 def timestamptz_to_micros(timestamptz_str: str) -> int:
     """Convert an ISO-8601 formatted timestamp with zone to microseconds from 1970-01-01T00:00:00.000000+00:00."""
-    if ISO_TIMESTAMPTZ.fullmatch(timestamptz_str):
+    if match := ISO_TIMESTAMPTZ.fullmatch(timestamptz_str):
+        offset = match.group(2)
+        timestamptz_str = timestamptz_str[: -len(offset)] + _utc_offset(offset)
         return datetime_to_micros(datetime.fromisoformat(timestamptz_str))
     if ISO_TIMESTAMP.fullmatch(timestamptz_str):
         # When we can match a timestamp without a zone, we can give a more specific error

tests/utils/test_datetime.py

@@ -29,6 +29,7 @@
     time_str_to_nanos,
     time_to_nanos,
     timestamp_to_nanos,
+    timestamptz_to_micros,
     timestamptz_to_nanos,
 )
 
@@ -139,12 +140,33 @@ def test_timestamp_to_nanos_invalid_timestamp() -> None:
         ("2025-02-23T16:21:44.3756127-04:00", 1740342104375612700),
         ("2025-02-23T16:21:44.37561278-04:00", 1740342104375612780),
         ("2025-02-23T16:21:44.375612001-04:00", 1740342104375612001),
+        ("2025-02-23T20:21:44.375612001Z", 1740342104375612001),
+        ("1970-01-01T00:00:00Z", 0),
+        ("1970-01-01T00:00:00z", 0),
+        ("2025-02-23T20:21Z", 1740342060000000000),
+        ("2025-02-23T16:21-04:00", 1740342060000000000),
     ],
 )
 def test_timestamptz_to_nanos(timestamp: str, nanos: int) -> None:
     assert nanos == timestamptz_to_nanos(timestamp)
 
 
+@pytest.mark.parametrize(
+    "timestamp, micros",
+    [
+        ("2025-02-23T16:21:44.375612-04:00", 1740342104375612),
+        ("2025-02-23T20:21:44.375612+00:00", 1740342104375612),
+        ("2025-02-23T20:21:44.375612Z", 1740342104375612),
+        ("1970-01-01T00:00:00Z", 0),
+        ("1970-01-01T00:00:00z", 0),
+        ("2025-02-23T20:21Z", 1740342060000000),
+        ("2025-02-23T16:21-04:00", 1740342060000000),
+    ],
+)
+def test_timestamptz_to_micros(timestamp: str, micros: int) -> None:
+    assert micros == timestamptz_to_micros(timestamp)
+
+
 def test_timestamptz_to_nanos_missing_zone_offset() -> None:
     with pytest.raises(ValueError, match="Missing zone offset: 2025-02-23T20:21:44.375612001"):
         timestamptz_to_nanos("2025-02-23T20:21:44.375612001")