DAS-2064: rename/relocate mapping.py -> xarray.core.datatree_mapping.py

DAS-2064: fix circular import issue.

xarray/core/datatree.py

@@ -18,6 +18,11 @@
 from xarray.core.coordinates import DatasetCoordinates
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset, DataVariables
+from xarray.core.datatree_mapping import (
+    TreeIsomorphismError,
+    check_isomorphic,
+    map_over_subtree,
+)
 from xarray.core.indexes import Index, Indexes
 from xarray.core.merge import dataset_update_method
 from xarray.core.options import OPTIONS as XR_OPTS
@@ -36,11 +41,6 @@
 from xarray.datatree_.datatree.formatting_html import (
     datatree_repr as datatree_repr_html,
 )
-from xarray.datatree_.datatree.mapping import (
-    TreeIsomorphismError,
-    check_isomorphic,
-    map_over_subtree,
-)
 from xarray.datatree_.datatree.ops import (
     DataTreeArithmeticMixin,
     MappedDatasetMethodsMixin,

xarray/datatree_/datatree/mapping.py → xarray/core/datatree_mapping.py

@@ -4,10 +4,9 @@
 import sys
 from itertools import repeat
 from textwrap import dedent
-from typing import TYPE_CHECKING, Callable, Tuple
+from typing import TYPE_CHECKING, Callable
 
 from xarray import DataArray, Dataset
-
 from xarray.core.iterators import LevelOrderIter
 from xarray.core.treenode import NodePath, TreeNode
 
@@ -154,7 +153,7 @@ def map_over_subtree(func: Callable) -> Callable:
     # TODO inspect function to work out immediately if the wrong number of arguments were passed for it?
 
     @functools.wraps(func)
-    def _map_over_subtree(*args, **kwargs) -> DataTree | Tuple[DataTree, ...]:
+    def _map_over_subtree(*args, **kwargs) -> DataTree | tuple[DataTree, ...]:
         """Internal function which maps func over every node in tree, returning a tree of the results."""
         from xarray.core.datatree import DataTree
 

xarray/datatree_/datatree/__init__.py

@@ -1,11 +1,8 @@
 # import public API
-from .mapping import TreeIsomorphismError, map_over_subtree
 from xarray.core.treenode import InvalidTreeError, NotFoundInTreeError
 
 
 __all__ = (
-    "TreeIsomorphismError",
     "InvalidTreeError",
     "NotFoundInTreeError",
-    "map_over_subtree",
 )

xarray/datatree_/datatree/formatting.py

@@ -2,7 +2,7 @@
 
 from xarray.core.formatting import _compat_to_str, diff_dataset_repr
 
-from xarray.datatree_.datatree.mapping import diff_treestructure
+from xarray.core.datatree_mapping import diff_treestructure
 from xarray.datatree_.datatree.render import RenderTree
 
 if TYPE_CHECKING:

xarray/datatree_/datatree/ops.py

@@ -2,7 +2,7 @@
 
 from xarray import Dataset
 
-from .mapping import map_over_subtree
+from xarray.core.datatree_mapping import map_over_subtree
 
 """
 Module which specifies the subset of xarray.Dataset's API which we wish to copy onto DataTree.

xarray/datatree_/datatree/tests/test_mapping.py → xarray/tests/test_datatree_mapping.py

@@ -1,9 +1,13 @@
 import numpy as np
 import pytest
-import xarray as xr
 
+import xarray as xr
 from xarray.core.datatree import DataTree
-from xarray.datatree_.datatree.mapping import TreeIsomorphismError, check_isomorphic, map_over_subtree
+from xarray.core.datatree_mapping import (
+    TreeIsomorphismError,
+    check_isomorphic,
+    map_over_subtree,
+)
 from xarray.datatree_.datatree.testing import assert_equal
 
 empty = xr.Dataset()
@@ -12,7 +16,7 @@
 class TestCheckTreesIsomorphic:
     def test_not_a_tree(self):
         with pytest.raises(TypeError, match="not a tree"):
-            check_isomorphic("s", 1)
+            check_isomorphic("s", 1)  # type: ignore[arg-type]
 
     def test_different_widths(self):
         dt1 = DataTree.from_dict(d={"a": empty})
@@ -69,7 +73,7 @@ def test_not_isomorphic_complex_tree(self, create_test_datatree):
     def test_checking_from_root(self, create_test_datatree):
         dt1 = create_test_datatree()
         dt2 = create_test_datatree()
-        real_root = DataTree(name="real root")
+        real_root: DataTree = DataTree(name="real root")
         dt2.name = "not_real_root"
         dt2.parent = real_root
         with pytest.raises(TreeIsomorphismError):