Inzynierka/Lib/site-packages/pandas/tests/generic/test_to_xarray.py

129 lines
4.0 KiB
Python
Raw Permalink Normal View History

2023-06-02 12:51:02 +02:00
import numpy as np
import pytest
import pandas.util._test_decorators as td
from pandas import (
Categorical,
DataFrame,
MultiIndex,
Series,
date_range,
)
import pandas._testing as tm
@td.skip_if_no("xarray")
class TestDataFrameToXArray:
@pytest.fixture
def df(self):
return DataFrame(
{
"a": list("abc"),
"b": list(range(1, 4)),
"c": np.arange(3, 6).astype("u1"),
"d": np.arange(4.0, 7.0, dtype="float64"),
"e": [True, False, True],
"f": Categorical(list("abc")),
"g": date_range("20130101", periods=3),
"h": date_range("20130101", periods=3, tz="US/Eastern"),
}
)
def test_to_xarray_index_types(self, index_flat, df):
index = index_flat
# MultiIndex is tested in test_to_xarray_with_multiindex
if len(index) == 0:
pytest.skip("Test doesn't make sense for empty index")
from xarray import Dataset
df.index = index[:3]
df.index.name = "foo"
df.columns.name = "bar"
result = df.to_xarray()
assert result.dims["foo"] == 3
assert len(result.coords) == 1
assert len(result.data_vars) == 8
tm.assert_almost_equal(list(result.coords.keys()), ["foo"])
assert isinstance(result, Dataset)
# idempotency
# datetimes w/tz are preserved
# column names are lost
expected = df.copy()
expected["f"] = expected["f"].astype(object)
expected.columns.name = None
tm.assert_frame_equal(result.to_dataframe(), expected)
def test_to_xarray_empty(self, df):
from xarray import Dataset
df.index.name = "foo"
result = df[0:0].to_xarray()
assert result.dims["foo"] == 0
assert isinstance(result, Dataset)
def test_to_xarray_with_multiindex(self, df):
from xarray import Dataset
# MultiIndex
df.index = MultiIndex.from_product([["a"], range(3)], names=["one", "two"])
result = df.to_xarray()
assert result.dims["one"] == 1
assert result.dims["two"] == 3
assert len(result.coords) == 2
assert len(result.data_vars) == 8
tm.assert_almost_equal(list(result.coords.keys()), ["one", "two"])
assert isinstance(result, Dataset)
result = result.to_dataframe()
expected = df.copy()
expected["f"] = expected["f"].astype(object)
expected.columns.name = None
tm.assert_frame_equal(result, expected)
@td.skip_if_no("xarray")
class TestSeriesToXArray:
def test_to_xarray_index_types(self, index_flat):
index = index_flat
# MultiIndex is tested in test_to_xarray_with_multiindex
from xarray import DataArray
ser = Series(range(len(index)), index=index, dtype="int64")
ser.index.name = "foo"
result = ser.to_xarray()
repr(result)
assert len(result) == len(index)
assert len(result.coords) == 1
tm.assert_almost_equal(list(result.coords.keys()), ["foo"])
assert isinstance(result, DataArray)
# idempotency
tm.assert_series_equal(result.to_series(), ser)
def test_to_xarray_empty(self):
from xarray import DataArray
ser = Series([], dtype=object)
ser.index.name = "foo"
result = ser.to_xarray()
assert len(result) == 0
assert len(result.coords) == 1
tm.assert_almost_equal(list(result.coords.keys()), ["foo"])
assert isinstance(result, DataArray)
def test_to_xarray_with_multiindex(self):
from xarray import DataArray
mi = MultiIndex.from_product([["a", "b"], range(3)], names=["one", "two"])
ser = Series(range(6), dtype="int64", index=mi)
result = ser.to_xarray()
assert len(result) == 2
tm.assert_almost_equal(list(result.coords.keys()), ["one", "two"])
assert isinstance(result, DataArray)
res = result.to_series()
tm.assert_series_equal(res, ser)