Skip to content

Commit a8c9896

Browse files
authored
Don't return IndexVariable with .dt accessor (#9415)
* Don't return IndexVariable with .dt accessor Closes #9190 * Support unstacking of index coordinates * Finish test * Fix mypy
1 parent a539b72 commit a8c9896

File tree

5 files changed

+34
-7
lines changed

5 files changed

+34
-7
lines changed

xarray/core/accessor_dt.py

+13-3
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
is_np_timedelta_like,
1515
)
1616
from xarray.core.types import T_DataArray
17-
from xarray.core.variable import IndexVariable
17+
from xarray.core.variable import IndexVariable, Variable
1818
from xarray.namedarray.utils import is_duck_dask_array
1919

2020
if TYPE_CHECKING:
@@ -244,12 +244,22 @@ def _date_field(self, name: str, dtype: DTypeLike) -> T_DataArray:
244244
if dtype is None:
245245
dtype = self._obj.dtype
246246
result = _get_date_field(_index_or_data(self._obj), name, dtype)
247-
newvar = self._obj.variable.copy(data=result, deep=False)
247+
newvar = Variable(
248+
dims=self._obj.dims,
249+
attrs=self._obj.attrs,
250+
encoding=self._obj.encoding,
251+
data=result,
252+
)
248253
return self._obj._replace(newvar, name=name)
249254

250255
def _tslib_round_accessor(self, name: str, freq: str) -> T_DataArray:
251256
result = _round_field(_index_or_data(self._obj), name, freq)
252-
newvar = self._obj.variable.copy(data=result, deep=False)
257+
newvar = Variable(
258+
dims=self._obj.dims,
259+
attrs=self._obj.attrs,
260+
encoding=self._obj.encoding,
261+
data=result,
262+
)
253263
return self._obj._replace(newvar, name=name)
254264

255265
def floor(self, freq: str) -> T_DataArray:

xarray/core/variable.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -1491,7 +1491,7 @@ def _unstack_once(
14911491
dim: Hashable,
14921492
fill_value=dtypes.NA,
14931493
sparse: bool = False,
1494-
) -> Self:
1494+
) -> Variable:
14951495
"""
14961496
Unstacks this variable given an index to unstack and the name of the
14971497
dimension to which the index refers.
@@ -1551,10 +1551,10 @@ def _unstack_once(
15511551
# case the destinations will be NaN / zero.
15521552
data[(..., *indexer)] = reordered
15531553

1554-
return self._replace(dims=new_dims, data=data)
1554+
return self.to_base_variable()._replace(dims=new_dims, data=data)
15551555

15561556
@partial(deprecate_dims, old_name="dimensions")
1557-
def unstack(self, dim=None, **dim_kwargs):
1557+
def unstack(self, dim=None, **dim_kwargs) -> Variable:
15581558
"""
15591559
Unstack an existing dimension into multiple new dimensions.
15601560

xarray/testing/assertions.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -211,7 +211,9 @@ def assert_identical(a, b, from_root=True):
211211
if isinstance(a, Variable):
212212
assert a.identical(b), formatting.diff_array_repr(a, b, "identical")
213213
elif isinstance(a, DataArray):
214-
assert a.name == b.name
214+
assert (
215+
a.name == b.name
216+
), f"DataArray names are different. L: {a.name}, R: {b.name}"
215217
assert a.identical(b), formatting.diff_array_repr(a, b, "identical")
216218
elif isinstance(a, Dataset | Variable):
217219
assert a.identical(b), formatting.diff_dataset_repr(a, b, "identical")

xarray/tests/test_accessor_dt.py

+1
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ def test_field_access(self, field) -> None:
9797
actual = getattr(self.data.time.dt, field)
9898
else:
9999
actual = getattr(self.data.time.dt, field)
100+
assert not isinstance(actual.variable, xr.IndexVariable)
100101

101102
assert expected.dtype == actual.dtype
102103
assert_identical(expected, actual)

xarray/tests/test_dataarray.py

+14
Original file line numberDiff line numberDiff line change
@@ -7202,3 +7202,17 @@ def test_lazy_data_variable_not_loaded():
72027202
da = xr.DataArray(v)
72037203
# No data needs to be accessed, so no error should be raised
72047204
xr.DataArray(da)
7205+
7206+
7207+
def test_unstack_index_var() -> None:
7208+
source = xr.DataArray(range(2), dims=["x"], coords=[["a", "b"]])
7209+
da = source.x
7210+
da = da.assign_coords(y=("x", ["c", "d"]), z=("x", ["e", "f"]))
7211+
da = da.set_index(x=["y", "z"])
7212+
actual = da.unstack("x")
7213+
expected = xr.DataArray(
7214+
np.array([["a", np.nan], [np.nan, "b"]], dtype=object),
7215+
coords={"y": ["c", "d"], "z": ["e", "f"]},
7216+
name="x",
7217+
)
7218+
assert_identical(actual, expected)

0 commit comments

Comments
 (0)