Skip to content

Commit

Permalink
use a composite strategy to generate the dataframe with a tz-aware …
Browse files Browse the repository at this point in the history
…datetime column (#9174)

* use a `composite` to generate the dataframe with a tz-aware dt column

* remove the `xfail`
  • Loading branch information
keewis authored Jul 1, 2024
1 parent 90e4486 commit 6c2d8c3
Showing 1 changed file with 23 additions and 16 deletions.
39 changes: 23 additions & 16 deletions properties/test_pandas_roundtrip.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
import pytest

import xarray as xr
from xarray.tests import has_pandas_3

pytest.importorskip("hypothesis")
import hypothesis.extra.numpy as npst # isort:skip
Expand All @@ -25,22 +24,34 @@

numeric_series = numeric_dtypes.flatmap(lambda dt: pdst.series(dtype=dt))


@st.composite
def dataframe_strategy(draw):
tz = draw(st.timezones())
dtype = pd.DatetimeTZDtype(unit="ns", tz=tz)

datetimes = st.datetimes(
min_value=pd.Timestamp("1677-09-21T00:12:43.145224193"),
max_value=pd.Timestamp("2262-04-11T23:47:16.854775807"),
timezones=st.just(tz),
)

df = pdst.data_frames(
[
pdst.column("datetime_col", elements=datetimes),
pdst.column("other_col", elements=st.integers()),
],
index=pdst.range_indexes(min_size=1, max_size=10),
)
return draw(df).astype({"datetime_col": dtype})


an_array = npst.arrays(
dtype=numeric_dtypes,
shape=npst.array_shapes(max_dims=2), # can only convert 1D/2D to pandas
)


datetime_with_tz_strategy = st.datetimes(timezones=st.timezones())
dataframe_strategy = pdst.data_frames(
[
pdst.column("datetime_col", elements=datetime_with_tz_strategy),
pdst.column("other_col", elements=st.integers()),
],
index=pdst.range_indexes(min_size=1, max_size=10),
)


@st.composite
def datasets_1d_vars(draw) -> xr.Dataset:
"""Generate datasets with only 1D variables
Expand Down Expand Up @@ -111,11 +122,7 @@ def test_roundtrip_pandas_dataframe(df) -> None:
xr.testing.assert_identical(arr, roundtripped.to_xarray())


@pytest.mark.skipif(
has_pandas_3,
reason="fails to roundtrip on pandas 3 (see https://github.com/pydata/xarray/issues/9098)",
)
@given(df=dataframe_strategy)
@given(df=dataframe_strategy())
def test_roundtrip_pandas_dataframe_datetime(df) -> None:
# Need to name the indexes, otherwise Xarray names them 'dim_0', 'dim_1'.
df.index.name = "rows"
Expand Down

0 comments on commit 6c2d8c3

Please sign in to comment.