Support timezone-aware datetime dtypes in hypothesis.extra.pandas

Passing a DatetimeTZDtype - e.g. "datetime64[ns, UTC]" - as the dtype
of series, indexes, or column now generates columns where every value
shares that single timezone, at the resolution given by the dtype's
unit.

Values are generated as naive datetime64[unit] (the canonical UTC int64
storage, including NaT) and the timezone is attached at the array level
with tz_localize("UTC").tz_convert(tz).  Constructing from
timezone-aware scalars would silently corrupt values near the int64
bounds, and Python datetime objects cannot express the coarser-unit
range at all.

For UTC and fixed-offset timezones the values cover the dtype's full
representable range.  Timezones whose UTC offset varies over time
resolve offsets through stdlib datetimes - as does pandas when
displaying values - so those are clamped to keep localized wall times
within years 1-9999.

Requires pandas >= 2.1: earlier versions silently coerce non-ns units
to nanoseconds, and pandas 2.0.x raises OutOfBoundsDatetime when
formatting values outside the ns-representable range.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TpEzQ8tW58JE94fBuFWgw5

Author: claude

hypothesis/RELEASE.rst

@@ -0,0 +1,19 @@
+RELEASE_TYPE: minor
+
+This release adds support for generating timezone-aware datetimes in
+:ref:`hypothesis.extra.pandas <hypothesis-pandas>`.  You can now pass a
+:class:`~pandas.DatetimeTZDtype` - such as ``"datetime64[ns, UTC]"`` - as the
+``dtype`` of a :func:`~hypothesis.extra.pandas.series`,
+:func:`~hypothesis.extra.pandas.indexes`, or
+:func:`~hypothesis.extra.pandas.column`, and every value will share that single
+timezone (the only arrangement pandas supports outside of the ``object``
+dtype).  The datetime resolution is taken from the dtype, so you can also
+generate e.g. ``"datetime64[us, UTC]"`` columns, and for UTC and other
+fixed-offset timezones the generated values cover the full range representable
+at that resolution - which for coarser units is far wider than the
+``datetime64[ns]`` bounds of roughly 1677-2262 (:issue:`4020`).
+
+Timezone-aware generation requires pandas >= 2.1: earlier versions silently
+coerce other resolutions to nanoseconds, or crash when displaying values
+outside the nanosecond-representable range.  Generated values include ``NaT``
+unless you pass an elements strategy which excludes it.

hypothesis/src/hypothesis/extra/pandas/impl.py

@@ -12,7 +12,7 @@
 from collections.abc import Sequence
 from copy import copy
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Any, Generic, Union
 
 import numpy as np
@@ -39,6 +39,20 @@
 except ImportError:
     IntegerDtype = ()
 
+try:
+    from pandas._libs.tslibs.timezones import is_fixed_offset, is_utc
+
+    def has_fixed_offset(tz):
+        return is_utc(tz) or is_fixed_offset(tz)
+
+except ImportError:
+
+    def has_fixed_offset(tz):
+        return isinstance(tz, timezone)
+
+
+PANDAS_GE_21 = tuple(int(x) for x in pandas.__version__.split(".")[:2]) >= (2, 1)
+
 
 def dtype_for_elements_strategy(s):
     return st.shared(
@@ -53,6 +67,74 @@ def infer_dtype_if_necessary(dtype, values, elements, draw):
     return dtype
 
 
+def datetime_tz_dtype(dtype):
+    """Return a normalised :class:`~pandas.DatetimeTZDtype` if ``dtype`` is a
+    timezone-aware datetime dtype (e.g. ``"datetime64[ns, UTC]"``), else None.
+
+    Every value in such a column or index shares a single timezone - the only
+    arrangement pandas supports outside of the ``object`` dtype.
+    """
+    tz_dtype = None
+    if isinstance(dtype, pandas.DatetimeTZDtype):
+        tz_dtype = dtype
+    elif isinstance(dtype, str):
+        try:
+            converted = pandas.api.types.pandas_dtype(dtype)
+        except TypeError:
+            return None
+        if isinstance(converted, pandas.DatetimeTZDtype):
+            tz_dtype = converted
+    if tz_dtype is not None and not PANDAS_GE_21:  # pragma: no cover
+        # Before pandas 2.0, non-nanosecond units silently coerce to ns; on
+        # pandas 2.0.x, formatting values outside the ns-representable range
+        # raises OutOfBoundsDatetime, so failing examples could not be shown.
+        raise InvalidArgument(
+            "Generating timezone-aware datetimes requires pandas >= 2.1, but "
+            f"you have pandas {pandas.__version__}"
+        )
+    return tz_dtype
+
+
+def naive_datetime64_dtype(tz_dtype):
+    """The underlying timezone-naive ``datetime64[unit]`` dtype that backs a
+    :class:`~pandas.DatetimeTZDtype`."""
+    return np.dtype(f"datetime64[{tz_dtype.unit}]")
+
+
+def tz_default_elements(tz_dtype):
+    """Default element strategy for a :class:`~pandas.DatetimeTZDtype`: naive
+    ``datetime64[unit]`` values, interpreted as UTC instants.
+
+    With a fixed-offset timezone (including UTC) these cover the dtype's full
+    representable range.  Other timezones resolve their UTC offsets through
+    stdlib datetimes - as does pandas when displaying values - so we keep each
+    instant a day inside Python's representable range, ensuring that the
+    localized wall times stay valid too.
+    """
+    naive = naive_datetime64_dtype(tz_dtype)
+    if has_fixed_offset(tz_dtype.tz):
+        return npst.from_dtype(naive)
+    unit = tz_dtype.unit
+    lo = int(np.datetime64(datetime.min + timedelta(days=1), unit).astype(np.int64))
+    hi = int(np.datetime64(datetime.max - timedelta(days=1), unit).astype(np.int64))
+    values = st.integers(lo, hi).map(lambda v: np.datetime64(v, unit))
+    return values | st.just(np.datetime64("NaT", unit))
+
+
+def attach_timezone(obj, tz):
+    """Interpret a naive datetime64 :class:`~pandas.Series` or
+    :class:`~pandas.Index` as UTC instants and convert it to ``tz``.
+
+    We build timezone-aware columns from naive ``datetime64[unit]`` values and
+    only attach the timezone at the array level, because constructing them from
+    timezone-aware scalars silently corrupts values near the bounds of the
+    representable range.
+    """
+    if isinstance(obj, pandas.Series):
+        return obj.dt.tz_localize("UTC").dt.tz_convert(tz)
+    return obj.tz_localize("UTC").tz_convert(tz)
+
+
 @check_function
 def elements_and_dtype(elements, dtype, source=None):
     if source is None:
@@ -103,6 +185,16 @@ def elements_and_dtype(elements, dtype, source=None):
             "here.  See https://stackoverflow.com/q/74355937 for workaround patterns."
         )
 
+    tz_dtype = datetime_tz_dtype(dtype)
+    if tz_dtype is not None:
+        # With the default elements strategy we build the column from the
+        # underlying naive datetime64[unit] dtype and the caller attaches the
+        # timezone afterwards.  Custom elements are passed through and
+        # constructed with the tz dtype.
+        if elements is None:
+            return tz_default_elements(tz_dtype), naive_datetime64_dtype(tz_dtype)
+        return elements, tz_dtype
+
     _get_subclasses = getattr(IntegerDtype, "__subclasses__", list)
     dtype = {t.name: t() for t in _get_subclasses()}.get(dtype, dtype)
 
@@ -246,11 +338,17 @@ def indexes(
     check_valid_interval(min_size, max_size, "min_size", "max_size")
     check_type(bool, unique, "unique")
 
+    tz_dtype = datetime_tz_dtype(dtype)
+    localize_tz = tz_dtype is not None and elements is None
+
     elements, dtype = elements_and_dtype(elements, dtype)
 
     if max_size is None:
         max_size = min_size + DEFAULT_MAX_SIZE
-    return ValueIndexStrategy(elements, dtype, min_size, max_size, unique, name)
+    strategy = ValueIndexStrategy(elements, dtype, min_size, max_size, unique, name)
+    if localize_tz:
+        return strategy.map(lambda ix: attach_timezone(ix, tz_dtype.tz))
+    return strategy
 
 
 @defines_strategy()
@@ -283,6 +381,15 @@ def series(
       elements strategy varies, then so will the resulting dtype of the
       series.
 
+      With pandas >= 2.1, you may also pass a timezone-aware
+      :class:`~pandas.DatetimeTZDtype` (e.g. ``"datetime64[ns, UTC]"``), in
+      which case every value in the series will share that single timezone.
+      For UTC and other fixed-offset timezones the generated values cover the
+      full range representable at the dtype's resolution - which for coarser
+      units is much wider than ``datetime64[ns]`` - while timezones with a
+      UTC offset that varies over time are limited to years 1-9999.  Values
+      include ``NaT`` unless you pass an elements strategy which excludes it.
+
     * index: If not None, a strategy for generating indexes for the
       resulting Series. This can generate either :class:`pandas.Index`
       objects or any sequence of values (which will be passed to the
@@ -309,11 +416,18 @@ def series(
     else:
         check_strategy(index, "index")
 
+    tz_dtype = datetime_tz_dtype(dtype)
+    localize_tz = tz_dtype is not None and elements is None
+
     elements, np_dtype = elements_and_dtype(elements, dtype)
     index_strategy = index
 
+    if localize_tz:
+        # Build the column with the underlying naive dtype, then attach the
+        # timezone at the array level once it has been assembled.
+        dtype = np_dtype
     # if it is converted to an object, use object for series type
-    if (
+    elif (
         np_dtype is not None
         and np_dtype.kind == "O"
         and not isinstance(dtype, IntegerDtype)
@@ -360,6 +474,8 @@ def result(draw):
                 name=draw(name),
             )
 
+    if localize_tz:
+        return result().map(lambda s: attach_timezone(s, tz_dtype.tz))
     return result()
 
 
@@ -561,6 +677,10 @@ def row():
 
     rewritten_columns = []
     column_names: set[str] = set()
+    # Maps the name of each timezone-aware datetime column to its timezone.  We
+    # build such columns with the underlying naive dtype and attach the
+    # timezone once the frame has been assembled (see attach_timezone).
+    tz_columns: dict = {}
 
     for i, c in enumerate(cols):
         check_type(column, c, f"columns[{i}]")
@@ -583,7 +703,12 @@ def row():
             raise InvalidArgument(f"duplicate definition of column name {c.name!r}")
 
         column_names.add(c.name)
-        c.elements, _ = elements_and_dtype(c.elements, c.dtype, label)
+        column_tz = datetime_tz_dtype(c.dtype)
+        localize_tz = column_tz is not None and c.elements is None
+        c.elements, np_dtype = elements_and_dtype(c.elements, c.dtype, label)
+        if localize_tz:
+            c.dtype = np_dtype
+            tz_columns[c.name] = column_tz.tz
 
         if c.dtype is None and rows is not None:
             raise InvalidArgument(
@@ -668,7 +793,10 @@ def just_draw_columns(draw):
                         )
                     )
 
-            return pandas.DataFrame(data, index=index)
+            result = pandas.DataFrame(data, index=index)
+            for name, tz in tz_columns.items():
+                result[name] = attach_timezone(result[name], tz)
+            return result
 
         return just_draw_columns()
     else:
@@ -756,6 +884,8 @@ def assign_rows(draw):
                     break
                 else:
                     reject()
+            for name, tz in tz_columns.items():
+                result[name] = attach_timezone(result[name], tz)
             return result
 
         return assign_rows()

hypothesis/tests/pandas/test_data_frame.py

@@ -15,7 +15,7 @@
 from hypothesis import HealthCheck, given, reject, settings, strategies as st
 from hypothesis.errors import InvalidArgument
 from hypothesis.extra import numpy as npst, pandas as pdst
-from hypothesis.extra.pandas.impl import IntegerDtype
+from hypothesis.extra.pandas.impl import PANDAS_GE_21, IntegerDtype
 
 from tests.common.debug import (
     assert_all_examples,
@@ -31,6 +31,48 @@ def test_can_have_columns_of_distinct_types(df):
     assert df["b"].dtype == np.dtype(float)
 
 
+requires_pandas21 = pytest.mark.skipif(
+    not PANDAS_GE_21, reason="timezone-aware dtypes require pandas >= 2.1"
+)
+
+
+@requires_pandas21
+def test_can_have_tz_aware_datetime_columns():
+    dtype = pd.DatetimeTZDtype(unit="us", tz="UTC")
+    assert_all_examples(
+        pdst.data_frames(
+            [pdst.column("a", dtype=dtype), pdst.column("b", dtype=dtype, unique=True)],
+            index=pdst.range_indexes(min_size=1),
+        ),
+        lambda df: df["a"].dtype == dtype
+        and df["b"].dtype == dtype
+        and df["b"].dropna().is_unique,
+    )
+
+
+@requires_pandas21
+def test_can_have_tz_aware_datetime_index():
+    dtype = pd.DatetimeTZDtype(unit="ns", tz="UTC")
+    assert_all_examples(
+        pdst.data_frames(
+            [pdst.column("a", dtype=int)], index=pdst.indexes(dtype=dtype, min_size=1)
+        ),
+        lambda df: df.index.dtype == dtype,
+    )
+
+
+@requires_pandas21
+def test_can_have_tz_aware_datetime_columns_with_rows():
+    dtype = pd.DatetimeTZDtype(unit="us", tz="UTC")
+    assert_all_examples(
+        pdst.data_frames(
+            [pdst.column("a", dtype=dtype)],
+            rows=st.tuples(st.datetimes()),
+        ),
+        lambda df: df["a"].dtype == dtype,
+    )
+
+
 @given(
     pdst.data_frames(
         [pdst.column(dtype=int)], index=pdst.range_indexes(min_size=1, max_size=5)

hypothesis/tests/pandas/test_indexes.py

@@ -17,8 +17,9 @@
 from hypothesis import HealthCheck, given, reject, settings, strategies as st
 from hypothesis.errors import Unsatisfiable
 from hypothesis.extra import numpy as npst, pandas as pdst
+from hypothesis.extra.pandas.impl import PANDAS_GE_21
 
-from tests.common.debug import check_can_generate_examples
+from tests.common.debug import assert_all_examples, check_can_generate_examples
 from tests.pandas.helpers import supported_by_pandas
 
 
@@ -63,6 +64,19 @@ def test_name_passed_on_indexes(s):
     assert s.name == "test_name"
 
 
+@pytest.mark.skipif(
+    not PANDAS_GE_21, reason="timezone-aware dtypes require pandas >= 2.1"
+)
+@pytest.mark.parametrize("unit", ["s", "ms", "us", "ns"])
+def test_tz_aware_datetime_indexes(unit):
+    dtype = pandas.DatetimeTZDtype(unit=unit, tz="UTC")
+    check_can_generate_examples(pdst.indexes(dtype=dtype, min_size=1))
+    assert_all_examples(
+        pdst.indexes(dtype=dtype, min_size=1, unique=True),
+        lambda ix: ix.dtype == dtype and ix.dropna().is_unique,
+    )
+
+
 # Sizes that fit into an int64 without overflow
 range_sizes = st.integers(0, 2**63 - 1)