Correctly handle zoneinfo timezones (#2253)

#### Reference Issues/PRs
Fixes #2245

#### What does this implement or fix?
Correctly normalizes zoneinfo timezones. Currently if we write something
as a pytz or a ZoneInfo timezone we'll read it back as a `pytz` to
preserve old pytz behavior and not break old arctic readers.


[Consensus](pandas-dev/pandas#34916 (comment))
is that people generally prefer `ZoneInfo` and we can change to always
using `ZoneInfo` but we'll need to do so in next major release.

#### Any other comments?

#### Checklist

<details>
  <summary>
   Checklist for code changes...
  </summary>
 
- [ ] Have you updated the relevant docstrings, documentation and
copyright notice?
- [ ] Is this contribution tested against [all ArcticDB's
features](../docs/mkdocs/docs/technical/contributing.md)?
- [ ] Do all exceptions introduced raise appropriate [error
messages](https://docs.arcticdb.io/error_messages/)?
 - [ ] Are API changes highlighted in the PR description?
- [ ] Is the PR labelled as enhancement or bug so it appears in
autogenerated release notes?
</details>

<!--
Thanks for contributing a Pull Request to ArcticDB! Please ensure you
have taken a look at:
- ArcticDB's Code of Conduct:
https://github.com/man-group/ArcticDB/blob/master/CODE_OF_CONDUCT.md
- ArcticDB's Contribution Licensing:
https://github.com/man-group/ArcticDB/blob/master/docs/mkdocs/docs/technical/contributing.md#contribution-licensing
-->

Author: IvoDD

python/arcticdb/version_store/_normalization.py

@@ -9,6 +9,9 @@
 import copy
 import datetime
 import io
+import sys
+if sys.version_info >= (3, 9):
+    import zoneinfo
 from datetime import timedelta
 import math
 
@@ -326,7 +329,7 @@ def _to_tz_timestamp(dt):
         second=dt.second,
         microsecond=dt.microsecond,
     ).value
-    tz = dt.tzinfo.zone if dt.tzinfo is not None else None
+    tz = _ensure_str_timezone(get_timezone(dt.tzinfo)) if dt.tzinfo is not None else None
     return [ts, tz]
 
 
@@ -400,14 +403,21 @@ def _normalize_single_index(
         return [str(name) for name in index_names], ix_vals
 
 
-def _ensure_str_timezone(index_tz):
-    if isinstance(index_tz, datetime.tzinfo) and check_is_utc_if_newer_pandas(index_tz):
+def _ensure_str_timezone(tz):
+    if sys.version_info >= (3, 9):
+        if isinstance(tz, zoneinfo.ZoneInfo):
+            # Pandas also special cases zoneinfo. It always just returns it as is instead of converting to string.
+            # Since we also want to support zoneinfo timezones and we want old arcticdb clients to read them,
+            # we convert them to a string (similarly to pytz timezones).
+            return str(tz)
+
+    if isinstance(tz, datetime.tzinfo) and check_is_utc_if_newer_pandas(tz):
         # Pandas started to treat UTC as a special case and give back the tzinfo object for it. We coerce it back to
         # a str to avoid special cases for it further along our pipeline. The breaking change was:
         # https://github.com/jbrockmendel/pandas/commit/94ce05d1bcc3c99e992c48cc99d0fd2726f43102#diff-3dba9e959e6ad7c394f0662a0e6477593fca446a6924437701ecff82b0b20b55
         return "UTC"
-    else:
-        return index_tz
+
+    return tz
 
 
 def _denormalize_single_index(item, norm_meta):
@@ -1019,7 +1029,7 @@ def denormalize(self, obj, meta):
 
     def _custom_pack(self, obj):
         if isinstance(obj, pd.Timestamp):
-            tz = obj.tz.zone if obj.tz is not None else None
+            tz = _ensure_str_timezone(get_timezone(obj.tz)) if obj.tz is not None else None
             return ExtType(MsgPackSerialization.PD_TIMESTAMP, packb([obj.value, tz]))
 
         if isinstance(obj, datetime.datetime):
@@ -1036,6 +1046,9 @@ def _custom_pack(self, obj):
     def _ext_hook(self, code, data):
         if code == MsgPackSerialization.PD_TIMESTAMP:
             data = unpackb(data, raw=False)
+            # TODO: Pandas by default interprets the `tz` string argument as `pytz`.
+            # We should instead use `ZoneInfo` but this will be an API break.
+            # See discussion here: https://github.com/pandas-dev/pandas/issues/34916#issuecomment-1595154742
             return pd.Timestamp(data[0], tz=data[1]) if data[1] is not None else pd.Timestamp(data[0])
 
         if code == MsgPackSerialization.PY_DATETIME:

python/tests/compat/arcticdb/test_compatibility.py

@@ -1,11 +1,15 @@
+import sys
 import pytest
+import pytz
+if sys.version_info >= (3, 9):
+    import zoneinfo
 from packaging import version
 import pandas as pd
 import numpy as np
 from arcticdb.util.test import assert_frame_equal
 from arcticdb.options import ModifiableEnterpriseLibraryOption
 from arcticdb.toolbox.library_tool import LibraryTool
-from tests.util.mark import ARCTICDB_USING_CONDA
+from tests.util.mark import ARCTICDB_USING_CONDA, ZONE_INFO_MARK
 from arcticdb_ext.tools import StorageMover
 
 from arcticdb.util.venv import CurrentVersion
@@ -187,6 +191,62 @@ def test_compat_snapshot_metadata_read(old_venv_and_arctic_uri, lib_name):
         assert meta == {"old_key": "old_value"}
 
 
+@ZONE_INFO_MARK
+@pytest.mark.parametrize("zone_name", ["UTC", "America/New_York"])
+def test_compat_timestamp_metadata(old_venv_and_arctic_uri, lib_name, zone_name):
+    old_venv, arctic_uri = old_venv_and_arctic_uri
+
+    old_ac = old_venv.create_arctic(arctic_uri)
+    old_lib = old_ac.create_library(lib_name)
+
+    sym = "sym"
+    df = pd.DataFrame({"col": [1]})
+    timestamp_no_tz = pd.Timestamp(2025, 1, 1)
+    timestamp_pytz = pd.Timestamp(2025, 1, 1, tz=pytz.timezone(zone_name))
+    timestamp_zoneinfo = pd.Timestamp(2025, 1, 1, tz=zoneinfo.ZoneInfo(zone_name))
+
+    # Write two versions with old arctic:
+    # v0 - no timezone
+    # v1 - pytz
+    old_lib.execute([
+        f"""
+import pytz
+timestamp_no_tz = pd.Timestamp(year=2025, month=1, day=1)
+timestamp_pytz = pd.Timestamp(year=2025, month=1, day=1, tz=pytz.timezone({repr(zone_name)}))
+lib.write({repr(sym)}, df, metadata=timestamp_no_tz)
+lib.write_metadata({repr(sym)}, timestamp_pytz)
+assert lib.read_metadata({repr(sym)}, as_of=0).metadata == timestamp_no_tz
+assert lib.read_metadata({repr(sym)}, as_of=1).metadata == timestamp_pytz
+        """
+    ], dfs={"df": df})
+
+    # Write 3 more versions with current arctic:
+    # v2 - no timezone
+    # v3 - pytz
+    # v4 - zoneinfo
+    with CurrentVersion(arctic_uri, lib_name) as curr:
+        lib = curr.lib
+        assert lib.read_metadata(sym, as_of=0).metadata == timestamp_no_tz
+        assert lib.read_metadata(sym, as_of=1).metadata == timestamp_pytz
+        lib.write_metadata(sym, timestamp_no_tz) # v2
+        lib.write_metadata(sym, timestamp_pytz) # v3
+        lib.write_metadata(sym, timestamp_zoneinfo) # v4
+        assert lib.read_metadata(sym, as_of=2).metadata == timestamp_no_tz
+        assert lib.read_metadata(sym, as_of=3).metadata == timestamp_pytz
+        assert lib.read_metadata(sym, as_of=4).metadata == timestamp_pytz
+
+    old_lib.execute([
+        f"""
+import pytz
+timestamp_no_tz = pd.Timestamp(year=2025, month=1, day=1)
+timestamp_pytz = pd.Timestamp(year=2025, month=1, day=1, tz=pytz.timezone({repr(zone_name)}))
+assert lib.read_metadata({repr(sym)}, as_of=2).metadata == timestamp_no_tz
+assert lib.read_metadata({repr(sym)}, as_of=3).metadata == timestamp_pytz
+assert lib.read_metadata({repr(sym)}, as_of=4).metadata == timestamp_pytz
+        """
+    ])
+
+
 def test_compat_read_incomplete(old_venv_and_arctic_uri, lib_name):
     old_venv, arctic_uri = old_venv_and_arctic_uri
     sym = "sym"

python/tests/integration/arcticdb/version_store/test_metadata_support.py

@@ -5,8 +5,15 @@
 
 As of the Change Date specified in that file, in accordance with the Business Source License, use of this software will be governed by the Apache License, version 2.0.
 """
+import sys
+
 import numpy as np
 import pandas as pd
+import datetime
+if sys.version_info >= (3, 9):
+    import zoneinfo
+import pytz
+
 from arcticdb_ext import set_config_string, unset_config_string
 from pandas import DataFrame, Timestamp
 import pytest
@@ -16,6 +23,8 @@
 from arcticdb_ext.storage import NoDataFoundException
 from arcticdb.util.test import assert_frame_equal, distinct_timestamps
 
+from tests.util.mark import ZONE_INFO_MARK
+
 
 # In the following lines, the naming convention is
 # test_rt_df stands for roundtrip dataframe (implicitly pandas given file name)
@@ -283,3 +292,30 @@ def test_rv_contains_metadata_batch_write_metadata(lmdb_version_store_v1):
     vits = lib.batch_write_metadata([sym_0, sym_1], [metadata_0, metadata_1])
     assert vits[0].metadata == metadata_0
     assert vits[1].metadata == metadata_1
+
+
+@ZONE_INFO_MARK
+@pytest.mark.parametrize("zone_type", ["pytz", "zoneinfo"])
+@pytest.mark.parametrize("zone_name", ["UTC", "America/Los_Angeles", "Europe/London", "Asia/Tokyo", "Pacific/Kiritimati"])
+def test_metadata_timestamp_with_tz(lmdb_version_store_v1, zone_type, zone_name):
+    lib = lmdb_version_store_v1
+    sym = "sym"
+    df = timestamp_indexed_df()
+    if zone_type == "pytz":
+        zone_to_write = pytz.timezone(zone_name)
+    elif zone_type == "zoneinfo":
+        zone_to_write = zoneinfo.ZoneInfo(zone_name)
+    else:
+        raise ValueError("Unknown timezone type")
+    # We need to normalize all timezoned timestamps to a pd.Timestamp with pytz
+    expected_metadata = pd.Timestamp(2025, 1, 1, tz=pytz.timezone(zone_name))
+
+    # pd.Timestamp
+    metadata_to_write = pd.Timestamp(2025, 1, 1, tz=zone_to_write)
+    lib.write(sym, df, metadata_to_write)
+    assert lib.read(sym).metadata == expected_metadata
+
+    # datetime.datetime
+    metadata_to_write = datetime.datetime(2025, 1, 1, tzinfo=zone_to_write)
+    lib.write(sym, df, metadata_to_write)
+    assert lib.read(sym).metadata == expected_metadata

python/tests/unit/arcticdb/version_store/test_normalization.py

@@ -12,11 +12,14 @@
 import sys
 from collections import namedtuple
 from unittest.mock import patch, MagicMock
+
 import numpy as np
 import pandas as pd
 import dateutil as du
 import pytest
 import pytz
+if sys.version_info >= (3, 9):
+    import zoneinfo
 from numpy.testing import assert_equal, assert_array_equal
 from arcticdb_ext.version_store import SortedValue as _SortedValue
 
@@ -49,7 +52,7 @@
 from arcticdb.util._versions import IS_PANDAS_ZERO, IS_PANDAS_TWO
 from arcticdb.exceptions import ArcticNativeException
 
-from tests.util.mark import param_dict
+from tests.util.mark import param_dict, ZONE_INFO_MARK
 
 params = {
     "simple_dict": {"a": "1", "b": 2, "c": 3.0, "d": True},
@@ -238,19 +241,35 @@ def test_empty_df():
     assert_frame_equal(d, D)
 
 
+timezone_params = [
+    "UTC",
+    "Europe/Amsterdam",
+    pytz.UTC,
+    pytz.timezone("Europe/Amsterdam"),
+    du.tz.gettz("UTC"),
+]
+if sys.version_info > (3, 9):
+    timezone_params += [
+        zoneinfo.ZoneInfo("UTC"),
+        zoneinfo.ZoneInfo("Pacific/Kiritimati"),
+        zoneinfo.ZoneInfo("America/Los_Angeles"),
+    ]
 # See test_get_description_date_range_tz in test_arctic.py for the V2 API equivalent
-@pytest.mark.parametrize(
-    "tz", ["UTC", "Europe/Amsterdam", pytz.UTC, pytz.timezone("Europe/Amsterdam"), du.tz.gettz("UTC")]
-)
+@pytest.mark.parametrize("tz", timezone_params)
 def test_write_tz(lmdb_version_store, sym, tz):
     assert tz is not None
-    index = index=pd.date_range(pd.Timestamp(0), periods=10, tz=tz)
+    index = pd.date_range(pd.Timestamp(0), periods=10, tz=tz)
     df = pd.DataFrame(data={"col1": np.arange(10)}, index=index)
     lmdb_version_store.write(sym, df)
     result = lmdb_version_store.read(sym).data
-    assert_frame_equal(df, result)
     df_tz = df.index.tzinfo
     assert str(df_tz) == str(tz)
+    expected_df = df
+    if sys.version_info >= (3, 9):
+        if isinstance(tz, zoneinfo.ZoneInfo):
+            # We convert all timezones to pytz, so we expect to convert to a pytz timezone.
+            expected_df.index = expected_df.index.tz_convert(pytz.timezone(str(tz)))
+    assert_frame_equal(expected_df, result)
     if tz == du.tz.gettz("UTC") and sys.version_info < (3, 7):
         pytest.skip("Timezone files don't seem to have ever worked properly on Python 3.6")
     start_ts, end_ts = lmdb_version_store.get_timerange_for_symbol(sym)

python/tests/util/mark.py

@@ -70,6 +70,10 @@
     not MEMRAY_SUPPORTED, reason="MEMRAY supports linux and macos and python 3.8 and above"
 )
 
+ZONE_INFO_MARK = pytest.mark.skipif(
+    sys.version_info < (3, 9),
+    reason="zoneinfo module was introduced in Python 3.9")
+
 SSL_TESTS_MARK = pytest.mark.skipif(
     not SSL_TEST_SUPPORTED,
     reason="When SSL tests are enabled",