Add broadcast_like.

xarray/core/alignment.py

@@ -391,6 +391,58 @@ def reindex_variables(
     return reindexed, new_indexes
 
 
+def _get_broadcast_dims_map_common_coords(args, exclude):
+
+    common_coords = OrderedDict()
+    dims_map = OrderedDict()
+    for arg in args:
+        for dim in arg.dims:
+            if dim not in common_coords and dim not in exclude:
+                dims_map[dim] = arg.sizes[dim]
+                if dim in arg.coords:
+                    common_coords[dim] = arg.coords[dim].variable
+
+    return dims_map, common_coords
+
+
+def _broadcast_helper(arg, exclude, dims_map, common_coords):
+
+    from .dataarray import DataArray
+    from .dataset import Dataset
+
+    def _set_dims(var):
+        # Add excluded dims to a copy of dims_map
+        var_dims_map = dims_map.copy()
+        for dim in exclude:
+            with suppress(ValueError):
+                # ignore dim not in var.dims
+                var_dims_map[dim] = var.shape[var.dims.index(dim)]
+
+        return var.set_dims(var_dims_map)
+
+    def _broadcast_array(array):
+        data = _set_dims(array.variable)
+        coords = OrderedDict(array.coords)
+        coords.update(common_coords)
+        return DataArray(data, coords, data.dims, name=array.name,
+                         attrs=array.attrs)
+
+    def _broadcast_dataset(ds):
+        data_vars = OrderedDict(
+            (k, _set_dims(ds.variables[k]))
+            for k in ds.data_vars)
+        coords = OrderedDict(ds.coords)
+        coords.update(common_coords)
+        return Dataset(data_vars, coords, ds.attrs)
+
+    if isinstance(arg, DataArray):
+        return _broadcast_array(arg)
+    elif isinstance(arg, Dataset):
+        return _broadcast_dataset(arg)
+    else:
+        raise ValueError('all input must be Dataset or DataArray objects')
+
+
 def broadcast(*args, exclude=None):
     """Explicitly broadcast any number of DataArray or Dataset objects against
     one another.
@@ -463,55 +515,16 @@ def broadcast(*args, exclude=None):
         a        (x, y) int64 1 1 2 2 3 3
         b        (x, y) int64 5 6 5 6 5 6
     """
-    from .dataarray import DataArray
-    from .dataset import Dataset
 
     if exclude is None:
         exclude = set()
     args = align(*args, join='outer', copy=False, exclude=exclude)
 
-    common_coords = OrderedDict()
-    dims_map = OrderedDict()
-    for arg in args:
-        for dim in arg.dims:
-            if dim not in common_coords and dim not in exclude:
-                dims_map[dim] = arg.sizes[dim]
-                if dim in arg.coords:
-                    common_coords[dim] = arg.coords[dim].variable
-
-    def _set_dims(var):
-        # Add excluded dims to a copy of dims_map
-        var_dims_map = dims_map.copy()
-        for dim in exclude:
-            with suppress(ValueError):
-                # ignore dim not in var.dims
-                var_dims_map[dim] = var.shape[var.dims.index(dim)]
-
-        return var.set_dims(var_dims_map)
-
-    def _broadcast_array(array):
-        data = _set_dims(array.variable)
-        coords = OrderedDict(array.coords)
-        coords.update(common_coords)
-        return DataArray(data, coords, data.dims, name=array.name,
-                         attrs=array.attrs)
-
-    def _broadcast_dataset(ds):
-        data_vars = OrderedDict(
-            (k, _set_dims(ds.variables[k]))
-            for k in ds.data_vars)
-        coords = OrderedDict(ds.coords)
-        coords.update(common_coords)
-        return Dataset(data_vars, coords, ds.attrs)
-
+    dims_map, common_coords = _get_broadcast_dims_map_common_coords(args,
+                                                                    exclude)
     result = []
     for arg in args:
-        if isinstance(arg, DataArray):
-            result.append(_broadcast_array(arg))
-        elif isinstance(arg, Dataset):
-            result.append(_broadcast_dataset(arg))
-        else:
-            raise ValueError('all input must be Dataset or DataArray objects')
+        result.append(_broadcast_helper(arg, exclude, dims_map, common_coords))
 
     return tuple(result)
 

xarray/core/dataarray.py

@@ -13,7 +13,9 @@
     computation, dtypes, groupby, indexing, ops, resample, rolling, utils)
 from .accessor_dt import DatetimeAccessor
 from .accessor_str import StringAccessor
-from .alignment import align, broadcast, reindex_like_indexers
+from .alignment import (align, _broadcast_helper,
+                        _get_broadcast_dims_map_common_coords,
+                        reindex_like_indexers)
 from .common import AbstractArray, DataWithCoords
 from .coordinates import (
     DataArrayCoordinates, LevelCoordinatesSource, assert_coordinate_consistent,
@@ -987,17 +989,27 @@ def sel_points(self, dim='points', method=None, tolerance=None,
         return self._from_temp_dataset(ds)
 
     def broadcast_like(self,
-                       other: Union['DataArray', Dataset]) -> 'DataArray':
+                       other: Union['DataArray', Dataset],
+                       exclude=None) -> 'DataArray':
         """Broadcast this DataArray against another Dataset or DataArray.
         This is equivalent to xr.broadcast(other, self)[1]
 
         Parameters
         ----------
         other : Dataset or DataArray
             Object against which to broadcast this array.
+        exclude : sequence of str, optional
+            Dimensions that must not be broadcasted
         """
 
-        return broadcast(other, self)[1]
+        if exclude is None:
+            exclude = set()
+        args = align(other, self, join='outer', copy=False, exclude=exclude)
+
+        dims_map, common_coords = _get_broadcast_dims_map_common_coords(
+            args, exclude)
+
+        return _broadcast_helper(self, exclude, dims_map, common_coords)
 
     def reindex_like(self, other: Union['DataArray', Dataset],
                      method: Optional[str] = None, tolerance=None,

xarray/core/dataset.py

@@ -15,7 +15,8 @@
 from ..coding.cftimeindex import _parse_array_of_cftime_strings
 from . import (alignment, dtypes, duck_array_ops, formatting, groupby,
                indexing, ops, pdcompat, resample, rolling, utils)
-from .alignment import align, broadcast
+from .alignment import (align, _broadcast_helper,
+                        _get_broadcast_dims_map_common_coords)
 from .common import (ALL_DIMS, DataWithCoords, ImplementsDatasetReduce,
                      _contains_datetime_like_objects)
 from .coordinates import (DatasetCoordinates, LevelCoordinatesSource,
@@ -1979,17 +1980,28 @@ def sel_points(self, dim='points', method=None, tolerance=None,
         return self.isel_points(dim=dim, **pos_indexers)
 
     def broadcast_like(self,
-                       other: Union['Dataset', 'DataArray']) -> 'Dataset':
+                       other: Union['Dataset', 'DataArray'],
+                       exclude=None) -> 'Dataset':
         """Broadcast this DataArray against another Dataset or DataArray.
         This is equivalent to xr.broadcast(other, self)[1]
 
         Parameters
         ----------
         other : Dataset or DataArray
             Object against which to broadcast this array.
+        exclude : sequence of str, optional
+            Dimensions that must not be broadcasted
+
         """
 
-        return broadcast(other, self)[1]
+        if exclude is None:
+            exclude = set()
+        args = align(other, self, join='outer', copy=False, exclude=exclude)
+
+        dims_map, common_coords = _get_broadcast_dims_map_common_coords(
+            args, exclude)
+
+        return _broadcast_helper(self, exclude, dims_map, common_coords)
 
     def reindex_like(self, other, method=None, tolerance=None, copy=True,
                      fill_value=dtypes.NA):