ensure other types defer to DataTree, thus fixing pydata#9365

TomNicholas · Oct 13, 2024 · 03ce2c5 · 03ce2c5
1 parent 909ae0e
commit 03ce2c5
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
@@ -4765,9 +4765,10 @@ def _unary_op(self, f: Callable, *args, **kwargs) -> Self:
     def _binary_op(
         self, other: DaCompatible, f: Callable, reflexive: bool = False
     ) -> Self:
+        from xarray.core.datatree import DataTree
         from xarray.core.groupby import GroupBy
 
-        if isinstance(other, Dataset | GroupBy):
+        if isinstance(other, DataTree | Dataset | GroupBy):
             return NotImplemented
         if isinstance(other, DataArray):
             align_type = OPTIONS["arithmetic_join"]

diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -7784,9 +7784,10 @@ def _unary_op(self, f, *args, **kwargs) -> Self:
 
     def _binary_op(self, other, f, reflexive=False, join=None) -> Dataset:
         from xarray.core.dataarray import DataArray
+        from xarray.core.datatree import DataTree
         from xarray.core.groupby import GroupBy
 
-        if isinstance(other, GroupBy):
+        if isinstance(other, DataTree | GroupBy):
             return NotImplemented
         align_type = OPTIONS["arithmetic_join"] if join is None else join
         if isinstance(other, DataArray | Dataset):

diff --git a/xarray/core/variable.py b/xarray/core/variable.py
@@ -2323,7 +2323,7 @@ def _unary_op(self, f, *args, **kwargs):
             return result
 
     def _binary_op(self, other, f, reflexive=False):
-        if isinstance(other, xr.DataArray | xr.Dataset):
+        if isinstance(other, xr.DataTree | xr.DataArray | xr.Dataset):
             return NotImplemented
         if reflexive and issubclass(type(self), type(other)):
             other_data, self_data, dims = _broadcast_compat_data(other, self)