MOD: Change DBNStore.to_df() timezone type

Author: nmacholl

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## 0.57.2 - TBD
+
+#### Enhancements
+- Changed the `tz` parameter in `DBNStore.to_df()` to accept `datetime.tzinfo` instead of `pytz.BaseTzInfo` explicitly
+
 ## 0.57.1 - 2025-06-17
 
 #### Enhancements

databento/common/dbnstore.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 import abc
+import datetime
 import decimal
 import itertools
 import logging
 import warnings
+import zoneinfo
 from collections.abc import Generator
 from collections.abc import Iterator
 from collections.abc import Mapping
@@ -28,7 +30,6 @@
 import pandas as pd
 import pyarrow as pa
 import pyarrow.parquet as pq
-import pytz
 import zstandard
 from databento_dbn import FIXED_PRICE_SCALE
 from databento_dbn import UNDEF_PRICE
@@ -859,7 +860,7 @@ def to_df(
         pretty_ts: bool = ...,
         map_symbols: bool = ...,
         schema: Schema | str | None = ...,
-        tz: pytz.BaseTzInfo | str = ...,
+        tz: datetime.tzinfo | str = ...,
         count: None = ...,
     ) -> pd.DataFrame: ...
 
@@ -870,7 +871,7 @@ def to_df(
         pretty_ts: bool = ...,
         map_symbols: bool = ...,
         schema: Schema | str | None = ...,
-        tz: pytz.BaseTzInfo | str = ...,
+        tz: datetime.tzinfo | str = ...,
         count: int = ...,
     ) -> DataFrameIterator: ...
 
@@ -880,8 +881,8 @@ def to_df(
         pretty_ts: bool = True,
         map_symbols: bool = True,
         schema: Schema | str | None = None,
-        tz: pytz.BaseTzInfo | str | Default[pytz.BaseTzInfo] = Default[pytz.BaseTzInfo](
-            pytz.UTC,
+        tz: datetime.tzinfo | str | Default[datetime.tzinfo] = Default[datetime.tzinfo](
+            datetime.timezone.utc,
         ),
         count: int | None = None,
     ) -> pd.DataFrame | DataFrameIterator:
@@ -909,7 +910,7 @@ def to_df(
         schema : Schema or str, optional
             The DBN schema for the dataframe.
             This is only required when reading a DBN stream with mixed record types.
-        tz : pytz.BaseTzInfo or str, default UTC
+        tz : datetime.tzinfo or str, default UTC
             If `pretty_ts` is `True`, all timestamps will be converted to the specified timezone.
         count : int, optional
             If set, instead of returning a single `DataFrame` a `DataFrameIterator`
@@ -939,8 +940,13 @@ def to_df(
                 "A timezone was specified when `pretty_ts` is `False`. Did you mean to set `pretty_ts=True`?",
             )
 
-        if not isinstance(tz, pytz.BaseTzInfo):
-            tz = pytz.timezone(tz)
+        if isinstance(tz, str):
+            tz = zoneinfo.ZoneInfo(tz)
+        elif not isinstance(tz, datetime.tzinfo):
+            raise ValueError(
+                f"The value {tz!r} is not a valid datetime.tzinfo",
+            )
+
         if schema is None:
             if self.schema is None:
                 raise ValueError("a schema must be specified for mixed DBN data")
@@ -1442,7 +1448,7 @@ def __init__(
         count: int | None,
         struct_type: type[DBNRecord],
         instrument_map: InstrumentMap,
-        tz: pytz.BaseTzInfo,
+        tz: datetime.tzinfo,
         price_type: PriceType = PriceType.FLOAT,
         pretty_ts: bool = True,
         map_symbols: bool = True,

tests/test_historical_bento.py

@@ -1,6 +1,7 @@
 import collections
 import datetime as dt
 import decimal
+import zoneinfo
 from io import BytesIO
 from pathlib import Path
 from typing import Any
@@ -1567,11 +1568,51 @@ def test_dbnstore_to_df_with_timezone(
     df.reset_index(inplace=True)
 
     # Assert
-    expected_timezone = pytz.timezone(timezone)._utcoffset
+    expected_timezone = zoneinfo.ZoneInfo(timezone).utcoffset(None)
     failures = []
     struct = SCHEMA_STRUCT_MAP[schema]
     for field in struct._timestamp_fields:
-        if df[field].dt.tz._utcoffset != expected_timezone:
+        if df[field].dt.tz.utcoffset(None) != expected_timezone:
+            failures.append(field)
+
+    assert not failures
+
+
+@pytest.mark.parametrize(
+    "timezone",
+    [
+        pytz.timezone("US/Central"),
+        pytz.timezone("US/Eastern"),
+        pytz.timezone("Europe/Vienna"),
+        pytz.timezone("Asia/Dubai"),
+        pytz.timezone("UTC"),
+    ],
+)
+@pytest.mark.parametrize(
+    "schema",
+    [pytest.param(schema, id=str(schema)) for schema in Schema.variants()],
+)
+def test_dbnstore_to_df_with_pytz_timezone(
+    test_data_path: Callable[[Dataset, Schema], Path],
+    schema: Schema,
+    timezone: pytz.BaseTzInfo,
+) -> None:
+    """
+    Test that setting the `tz` parameter in `DBNStore.to_df` accepts `pytz`
+    timezone objects.
+    """
+    # Arrange
+    dbnstore = DBNStore.from_file(path=test_data_path(Dataset.GLBX_MDP3, schema))
+
+    # Act
+    df = dbnstore.to_df(tz=timezone)
+    df.reset_index(inplace=True)
+
+    # Assert
+    failures = []
+    struct = SCHEMA_STRUCT_MAP[schema]
+    for field in struct._timestamp_fields:
+        if df[field].dt.tz != timezone:
             failures.append(field)
 
     assert not failures
@@ -1591,7 +1632,7 @@ def test_dbnstore_to_df_with_timezone_pretty_ts_error(
     with pytest.raises(ValueError):
         dbnstore.to_df(
             pretty_ts=False,
-            tz=pytz.UTC,
+            tz=zoneinfo.ZoneInfo("UTC"),
         )