[namedarray] split .set_dims() into .expand_dims() and broadcast_to()

xarray/core/dataarray.py

@@ -4,15 +4,7 @@
 import warnings
 from collections.abc import Hashable, Iterable, Mapping, MutableMapping, Sequence
 from os import PathLike
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Generic,
-    Literal,
-    NoReturn,
-    overload,
-)
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, NoReturn, overload
 
 import numpy as np
 import pandas as pd
@@ -64,6 +56,7 @@
     as_compatible_data,
     as_variable,
 )
+from xarray.namedarray.utils import infix_dims
 from xarray.plot.accessor import DataArrayPlotAccessor
 from xarray.plot.utils import _get_units_from_attrs
 from xarray.util.deprecation_helpers import _deprecate_positional_args
@@ -2994,7 +2987,7 @@ def transpose(
         Dataset.transpose
         """
         if dims:
-            dims = tuple(utils.infix_dims(dims, self.dims, missing_dims))
+            dims = tuple(infix_dims(dims, self.dims, missing_dims))
         variable = self.variable.transpose(*dims)
         if transpose_coords:
             coords: dict[Hashable, Variable] = {}

xarray/core/dataset.py

@@ -110,8 +110,6 @@
     decode_numpy_dict_values,
     drop_dims_from_indexers,
     either_dict_or_kwargs,
-    infix_dims,
-    is_dict_like,
     is_scalar,
     maybe_wrap_array,
 )
@@ -122,6 +120,7 @@
     broadcast_variables,
     calculate_dimensions,
 )
+from xarray.namedarray.utils import infix_dims, is_dict_like
 from xarray.plot.accessor import DatasetPlotAccessor
 from xarray.util.deprecation_helpers import _deprecate_positional_args
 

xarray/core/utils.py

@@ -789,36 +789,6 @@ def __len__(self) -> int:
         return len(self._data) - num_hidden
 
 
-def infix_dims(
-    dims_supplied: Collection,
-    dims_all: Collection,
-    missing_dims: ErrorOptionsWithWarn = "raise",
-) -> Iterator:
-    """
-    Resolves a supplied list containing an ellipsis representing other items, to
-    a generator with the 'realized' list of all items
-    """
-    if ... in dims_supplied:
-        if len(set(dims_all)) != len(dims_all):
-            raise ValueError("Cannot use ellipsis with repeated dims")
-        if list(dims_supplied).count(...) > 1:
-            raise ValueError("More than one ellipsis supplied")
-        other_dims = [d for d in dims_all if d not in dims_supplied]
-        existing_dims = drop_missing_dims(dims_supplied, dims_all, missing_dims)
-        for d in existing_dims:
-            if d is ...:
-                yield from other_dims
-            else:
-                yield d
-    else:
-        existing_dims = drop_missing_dims(dims_supplied, dims_all, missing_dims)
-        if set(existing_dims) ^ set(dims_all):
-            raise ValueError(
-                f"{dims_supplied} must be a permuted list of {dims_all}, unless `...` is included"
-            )
-        yield from existing_dims
-
-
 def get_temp_dimname(dims: Container[Hashable], new_dim: Hashable) -> Hashable:
     """Get an new dimension name based on new_dim, that is not used in dims.
     If the same name exists, we add an underscore(s) in the head.

xarray/core/variable.py

@@ -42,7 +42,6 @@
     drop_dims_from_indexers,
     either_dict_or_kwargs,
     ensure_us_time_resolution,
-    infix_dims,
     is_duck_array,
     maybe_coerce_to_str,
 )
@@ -1382,107 +1381,6 @@ def roll(self, shifts=None, **shifts_kwargs):
             result = result._roll_one_dim(dim, count)
         return result
 
-    def transpose(
-        self,
-        *dims: Hashable | ellipsis,
-        missing_dims: ErrorOptionsWithWarn = "raise",
-    ) -> Self:
-        """Return a new Variable object with transposed dimensions.
-
-        Parameters
-        ----------
-        *dims : Hashable, optional
-            By default, reverse the dimensions. Otherwise, reorder the
-            dimensions to this order.
-        missing_dims : {"raise", "warn", "ignore"}, default: "raise"
-            What to do if dimensions that should be selected from are not present in the
-            Variable:
-            - "raise": raise an exception
-            - "warn": raise a warning, and ignore the missing dimensions
-            - "ignore": ignore the missing dimensions
-
-        Returns
-        -------
-        transposed : Variable
-            The returned object has transposed data and dimensions with the
-            same attributes as the original.
-
-        Notes
-        -----
-        This operation returns a view of this variable's data. It is
-        lazy for dask-backed Variables but not for numpy-backed Variables.
-
-        See Also
-        --------
-        numpy.transpose
-        """
-        if len(dims) == 0:
-            dims = self.dims[::-1]
-        else:
-            dims = tuple(infix_dims(dims, self.dims, missing_dims))
-
-        if len(dims) < 2 or dims == self.dims:
-            # no need to transpose if only one dimension
-            # or dims are in same order
-            return self.copy(deep=False)
-
-        axes = self.get_axis_num(dims)
-        data = as_indexable(self._data).transpose(axes)
-        return self._replace(dims=dims, data=data)
-
-    @property
-    def T(self) -> Self:
-        return self.transpose()
-
-    def set_dims(self, dims, shape=None):
-        """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 variable. If a dict, values are
-            used to provide the sizes of new dimensions; otherwise, new
-            dimensions are inserted with length 1.
-
-        Returns
-        -------
-        Variable
-        """
-        if isinstance(dims, str):
-            dims = [dims]
-
-        if shape is None and utils.is_dict_like(dims):
-            shape = dims.values()
-
-        missing_dims = set(self.dims) - set(dims)
-        if missing_dims:
-            raise ValueError(
-                f"new dimensions {dims!r} must be a superset of "
-                f"existing dimensions {self.dims!r}"
-            )
-
-        self_dims = set(self.dims)
-        expanded_dims = tuple(d for d in dims if d not in self_dims) + self.dims
-
-        if self.dims == expanded_dims:
-            # don't use broadcast_to unless necessary so the result remains
-            # writeable if possible
-            expanded_data = self.data
-        elif shape is not None:
-            dims_map = dict(zip(dims, shape))
-            tmp_shape = tuple(dims_map[d] for d in expanded_dims)
-            expanded_data = duck_array_ops.broadcast_to(self.data, tmp_shape)
-        else:
-            expanded_data = self.data[(None,) * (len(expanded_dims) - self.ndim)]
-
-        expanded_var = Variable(
-            expanded_dims, expanded_data, self._attrs, self._encoding, fastpath=True
-        )
-        return expanded_var.transpose(*dims)
-
     def _stack_once(self, dims: list[Hashable], new_dim: Hashable):
         if not set(dims) <= set(self.dims):
             raise ValueError(f"invalid existing dimensions: {dims}")

xarray/namedarray/_typing.py

@@ -6,6 +6,7 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    Literal,
     Protocol,
     SupportsIndex,
     TypeVar,
@@ -263,3 +264,6 @@ def todense(self) -> NDArray[_ScalarType_co]:
 
 # NamedArray can most likely use both __array_function__ and __array_namespace__:
 _sparsearrayfunction_or_api = (_sparsearrayfunction, _sparsearrayapi)
+
+ErrorOptions = Literal["raise", "ignore"]
+ErrorOptionsWithWarn = Literal["raise", "warn", "ignore"]

xarray/namedarray/core.py

@@ -22,17 +22,25 @@
 from xarray.core import dtypes, formatting, formatting_html
 from xarray.namedarray._aggregations import NamedArrayAggregations
 from xarray.namedarray._typing import (
+    ErrorOptionsWithWarn,
     _arrayapi,
     _arrayfunction_or_api,
     _chunkedarray,
     _dtype,
     _DType_co,
     _ScalarType_co,
+    _ShapeLike,
     _ShapeType_co,
     _SupportsImag,
     _SupportsReal,
 )
-from xarray.namedarray.utils import _default, is_duck_dask_array, to_0d_object_array
+from xarray.namedarray.utils import (
+    _default,
+    infix_dims,
+    is_dict_like,
+    is_duck_dask_array,
+    to_0d_object_array,
+)
 
 if TYPE_CHECKING:
     from numpy.typing import ArrayLike, NDArray
@@ -847,5 +855,108 @@ def _to_dense(self) -> Self:
         else:
             raise TypeError("self.data is not a sparse array")
 
+    def transpose(
+        self,
+        *dims: Hashable | ellipsis,
+        missing_dims: ErrorOptionsWithWarn = "raise",
+    ) -> Self:
+        """Return a new object with transposed dimensions.
+
+        Parameters
+        ----------
+        *dims : Hashable, optional
+            By default, reverse the dimensions. Otherwise, reorder the
+            dimensions to this order.
+        missing_dims : {"raise", "warn", "ignore"}, default: "raise"
+            What to do if dimensions that should be selected from are not present in the
+            NamedArray:
+            - "raise": raise an exception
+            - "warn": raise a warning, and ignore the missing dimensions
+            - "ignore": ignore the missing dimensions
+
+        Returns
+        -------
+        transposed : NamedArray
+            The returned object has transposed data and dimensions with the
+            same attributes as the original.
+
+        Notes
+        -----
+        This operation returns a view of this variable's data. It is
+        lazy for dask-backed Variables but not for numpy-backed Variables.
+
+        See Also
+        --------
+        numpy.transpose
+        """
+        from xarray.core.indexing import as_indexable
+
+        if len(dims) == 0:
+            dims = self.dims[::-1]
+        else:
+            dims = tuple(infix_dims(dims, self.dims, missing_dims))
+
+        if len(dims) < 2 or dims == self.dims:
+            # no need to transpose if only one dimension
+            # or dims are in same order
+            return self.copy(deep=False)
+
+        axes = self.get_axis_num(dims)
+        data = as_indexable(self._data).transpose(axes)
+        return self._replace(dims=dims, data=data)
+
+    @property
+    def T(self) -> Self:
+        return self.transpose()
+
+    def set_dims(self, dims: _DimsLike, shape: _ShapeLike = None) -> Self:
+        """
+        Return a new namedarray with given set of dimensions.
+        This method might be used to attach new dimension(s) to namedarray.
+
+        When possible, this operation does not copy this namedarray's data.
+
+        Parameters
+        ----------
+        dims : str or sequence of str or dict
+            Dimensions to include on the new namedarray. 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 of the new namedarray. If not provided, the shape is inferred from the data.
+        """
+
+        from xarray.core import duck_array_ops  # TODO: remove this import
+
+        if isinstance(dims, str):
+            dims = [dims]
+
+        if shape is None and is_dict_like(dims):
+            shape = dims.values()
+
+        missing_dims = set(self.dims) - set(dims)
+        if missing_dims:
+            raise ValueError(
+                f"new dimensions {dims!r} must be a superset of "
+                f"existing dimensions {self.dims!r}"
+            )
+
+        self_dims = set(self.dims)
+        expanded_dims = tuple(dim for dim in dims if dim not in self_dims) + self.dims
+
+        if self.dims == expanded_dims:
+            # don't use broadcast_to unless necessary so the result remains
+            # writeable if possible
+            expanded_data = self.data
+        elif shape is not None:
+            dims_map = dict(zip(dims, shape))
+            temporary_shape = tuple(dims_map[dim] for dim in expanded_dims)
+            expanded_data = duck_array_ops.broadcast_to(self.data, temporary_shape)
+        else:
+            expanded_data = self.data[(None,) * (len(expanded_dims) - self.ndim)]
+
+        expanded_obj = self._replace(data=expanded_data, dims=expanded_dims)
+        return expanded_obj.transpose(*dims)
+
 
 _NamedArray = NamedArray[Any, np.dtype[_ScalarType_co]]