Skip to content

Commit

Permalink
revert to .set_dims on Variable
Browse files Browse the repository at this point in the history
  • Loading branch information
andersy005 committed Oct 27, 2023
1 parent 8091265 commit 2b93273
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 14 deletions.
20 changes: 20 additions & 0 deletions xarray/core/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
is_duck_array,
maybe_coerce_to_str,
)
from xarray.namedarray._typing import _DimsLike, _ShapeLike
from xarray.namedarray.core import NamedArray

NON_NUMPY_SUPPORTED_ARRAY_TYPES = (
Expand Down Expand Up @@ -1381,6 +1382,25 @@ def roll(self, shifts=None, **shifts_kwargs):
result = result._roll_one_dim(dim, count)
return result

def set_dims(self, dims: _DimsLike, shape: _ShapeLike | None = None) -> Self:
"""
Return a new variable with given set of dimensions.
This method might be used to attach new dimension(s) to variable.
When possible, this operation does not copy this variable's data.
Parameters
----------
dims : str or sequence of str or dict
Dimensions to include on the new object (must be a superset of the existing dimensions).
If a dict, values are used to provide the sizes of new dimensions; otherwise, new dimensions are inserted with length 1.
shape : sequence of int, optional
Shape to broadcast the data to. Must be specified in the same order as `dims`.
If not provided, new dimensions are inserted with length 1.
"""

return self.expand_dims(dims, shape)

def _stack_once(self, dims: list[Hashable], new_dim: Hashable):
if not set(dims) <= set(self.dims):
raise ValueError(f"invalid existing dimensions: {dims}")
Expand Down
3 changes: 1 addition & 2 deletions xarray/namedarray/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -913,8 +913,7 @@ def T(self) -> Self:
def broadcast_to(self, shape: _ShapeLike) -> NamedArray[Any, _DType_co]:
from xarray.core import duck_array_ops # TODO: remove this import

data = duck_array_ops.broadcast_to(self.data, shape)
return self._replace(data=data)
return duck_array_ops.broadcast_to(self.data, shape)

def _create_expanded_obj(
self, expanded_data: duckarray[Any, Any], expanded_dims: _DimsLike
Expand Down
22 changes: 10 additions & 12 deletions xarray/namedarray/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@

import sys
import warnings
from collections.abc import Collection, Hashable, Iterable, Iterator, Mapping
from collections.abc import Hashable, Iterator, Mapping
from enum import Enum
from typing import TYPE_CHECKING, Any, Final

import numpy as np

from xarray.namedarray._typing import ErrorOptionsWithWarn
from xarray.namedarray._typing import ErrorOptionsWithWarn, _DimsLike

if TYPE_CHECKING:
if sys.version_info >= (3, 10):
Expand Down Expand Up @@ -82,10 +82,10 @@ def is_dict_like(value: Any) -> TypeGuard[Mapping[Any, Any]]:


def drop_missing_dims(
supplied_dims: Iterable[Hashable],
dims: Iterable[Hashable],
supplied_dims: _DimsLike,
dims: _DimsLike,
missing_dims: ErrorOptionsWithWarn,
) -> Iterable[Hashable]:
) -> _DimsLike:
"""Depending on the setting of missing_dims, drop any dimensions from supplied_dims that
are not present in dims.
Expand All @@ -98,17 +98,15 @@ def drop_missing_dims(

if missing_dims == "raise":
supplied_dims_set = {val for val in supplied_dims if val is not ...}
invalid = supplied_dims_set - set(dims)
if invalid:
if invalid := supplied_dims_set - set(dims):
raise ValueError(
f"Dimensions {invalid} do not exist. Expected one or more of {dims}"
)

return supplied_dims

elif missing_dims == "warn":
invalid = set(supplied_dims) - set(dims)
if invalid:
if invalid := set(supplied_dims) - set(dims):
warnings.warn(
f"Dimensions {invalid} do not exist. Expected one or more of {dims}"
)
Expand All @@ -125,10 +123,10 @@ def drop_missing_dims(


def infix_dims(
dims_supplied: Collection[Any],
dims_all: Collection[Any],
dims_supplied: _DimsLike,
dims_all: _DimsLike,
missing_dims: ErrorOptionsWithWarn = "raise",
) -> Iterator[Any]:
) -> Iterator[_DimsLike]:
"""
Resolves a supplied list containing an ellipsis representing other items, to
a generator with the 'realized' list of all items
Expand Down

0 comments on commit 2b93273

Please sign in to comment.