Move parallelcompat and chunkmanagers to NamedArray (#8319)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Anderson Banihirwe <[email protected]>
Co-authored-by: Anderson Banihirwe <[email protected]>

Author: 3 people

doc/internals/chunked-arrays.rst

@@ -35,44 +35,44 @@ The implementation of these functions is specific to the type of arrays passed t
 whereas :py:class:`cubed.Array` objects must be processed by :py:func:`cubed.map_blocks`.
 
 In order to use the correct implementation of a core operation for the array type encountered, xarray dispatches to the
-corresponding subclass of :py:class:`~xarray.core.parallelcompat.ChunkManagerEntrypoint`,
+corresponding subclass of :py:class:`~xarray.namedarray.parallelcompat.ChunkManagerEntrypoint`,
 also known as a "Chunk Manager". Therefore **a full list of the operations that need to be defined is set by the
-API of the** :py:class:`~xarray.core.parallelcompat.ChunkManagerEntrypoint` **abstract base class**. Note that chunked array
+API of the** :py:class:`~xarray.namedarray.parallelcompat.ChunkManagerEntrypoint` **abstract base class**. Note that chunked array
 methods are also currently dispatched using this class.
 
 Chunked array creation is also handled by this class. As chunked array objects have a one-to-one correspondence with
 in-memory numpy arrays, it should be possible to create a chunked array from a numpy array by passing the desired
-chunking pattern to an implementation of :py:class:`~xarray.core.parallelcompat.ChunkManagerEntrypoint.from_array``.
+chunking pattern to an implementation of :py:class:`~xarray.namedarray.parallelcompat.ChunkManagerEntrypoint.from_array``.
 
 .. note::
 
-    The :py:class:`~xarray.core.parallelcompat.ChunkManagerEntrypoint` abstract base class is mostly just acting as a
+    The :py:class:`~xarray.namedarray.parallelcompat.ChunkManagerEntrypoint` abstract base class is mostly just acting as a
     namespace for containing the chunked-aware function primitives. Ideally in the future we would have an API standard
     for chunked array types which codified this structure, making the entrypoint system unnecessary.
 
-.. currentmodule:: xarray.core.parallelcompat
+.. currentmodule:: xarray.namedarray.parallelcompat
 
-.. autoclass:: xarray.core.parallelcompat.ChunkManagerEntrypoint
+.. autoclass:: xarray.namedarray.parallelcompat.ChunkManagerEntrypoint
    :members:
 
 Registering a new ChunkManagerEntrypoint subclass
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 Rather than hard-coding various chunk managers to deal with specific chunked array implementations, xarray uses an
 entrypoint system to allow developers of new chunked array implementations to register their corresponding subclass of
-:py:class:`~xarray.core.parallelcompat.ChunkManagerEntrypoint`.
+:py:class:`~xarray.namedarray.parallelcompat.ChunkManagerEntrypoint`.
 
 
 To register a new entrypoint you need to add an entry to the ``setup.cfg`` like this::
 
     [options.entry_points]
     xarray.chunkmanagers =
-        dask = xarray.core.daskmanager:DaskManager
+        dask = xarray.namedarray.daskmanager:DaskManager
 
 See also `cubed-xarray <https://github.com/xarray-contrib/cubed-xarray>`_ for another example.
 
 To check that the entrypoint has worked correctly, you may find it useful to display the available chunkmanagers using
-the internal function :py:func:`~xarray.core.parallelcompat.list_chunkmanagers`.
+the internal function :py:func:`~xarray.namedarray.parallelcompat.list_chunkmanagers`.
 
 .. autofunction:: list_chunkmanagers
 

doc/whats-new.rst

@@ -88,6 +88,9 @@ Internal Changes
   when the data isn't datetime-like. (:issue:`8718`, :pull:`8724`)
   By `Maximilian Roos <https://github.com/max-sixty>`_.
 
+- Move `parallelcompat` and `chunk managers` modules from `xarray/core` to `xarray/namedarray`. (:pull:`8319`)
+  By `Tom Nicholas <https://github.com/TomNicholas>`_ and `Anderson Banihirwe <https://github.com/andersy005>`_.
+
 .. _whats-new.2024.01.1:
 
 v2024.01.1 (23 Jan, 2024)

pyproject.toml

@@ -54,7 +54,7 @@ issue-tracker = "https://github.com/pydata/xarray/issues"
 source-code = "https://github.com/pydata/xarray"
 
 [project.entry-points."xarray.chunkmanagers"]
-dask = "xarray.core.daskmanager:DaskManager"
+dask = "xarray.namedarray.daskmanager:DaskManager"
 
 [build-system]
 build-backend = "setuptools.build_meta"

xarray/backends/api.py

@@ -34,13 +34,13 @@
     _nested_combine,
     combine_by_coords,
 )
-from xarray.core.daskmanager import DaskManager
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset, _get_chunk, _maybe_chunk
 from xarray.core.indexes import Index
-from xarray.core.parallelcompat import guess_chunkmanager
 from xarray.core.types import ZarrWriteModes
 from xarray.core.utils import is_remote_uri
+from xarray.namedarray.daskmanager import DaskManager
+from xarray.namedarray.parallelcompat import guess_chunkmanager
 
 if TYPE_CHECKING:
     try:

xarray/backends/common.py

@@ -12,9 +12,9 @@
 
 from xarray.conventions import cf_encoder
 from xarray.core import indexing
-from xarray.core.parallelcompat import get_chunked_array_type
-from xarray.core.pycompat import is_chunked_array
 from xarray.core.utils import FrozenDict, NdimSizeLenMixin, is_remote_uri
+from xarray.namedarray.parallelcompat import get_chunked_array_type
+from xarray.namedarray.pycompat import is_chunked_array
 
 if TYPE_CHECKING:
     from io import BufferedIOBase

xarray/backends/pydap_.py

@@ -14,7 +14,6 @@
 )
 from xarray.backends.store import StoreBackendEntrypoint
 from xarray.core import indexing
-from xarray.core.pycompat import integer_types
 from xarray.core.utils import (
     Frozen,
     FrozenDict,
@@ -23,6 +22,7 @@
     is_remote_uri,
 )
 from xarray.core.variable import Variable
+from xarray.namedarray.pycompat import integer_types
 
 if TYPE_CHECKING:
     import os

xarray/backends/zarr.py

@@ -19,15 +19,15 @@
 )
 from xarray.backends.store import StoreBackendEntrypoint
 from xarray.core import indexing
-from xarray.core.parallelcompat import guess_chunkmanager
-from xarray.core.pycompat import integer_types
 from xarray.core.types import ZarrWriteModes
 from xarray.core.utils import (
     FrozenDict,
     HiddenKeyDict,
     close_on_error,
 )
 from xarray.core.variable import Variable
+from xarray.namedarray.parallelcompat import guess_chunkmanager
+from xarray.namedarray.pycompat import integer_types
 
 if TYPE_CHECKING:
     from io import BufferedIOBase

xarray/coding/strings.py

@@ -15,8 +15,9 @@
     unpack_for_encoding,
 )
 from xarray.core import indexing
-from xarray.core.parallelcompat import get_chunked_array_type, is_chunked_array
 from xarray.core.variable import Variable
+from xarray.namedarray.parallelcompat import get_chunked_array_type
+from xarray.namedarray.pycompat import is_chunked_array
 
 
 def create_vlen_dtype(element_type):

xarray/coding/times.py

@@ -24,11 +24,12 @@
 from xarray.core.common import contains_cftime_datetimes, is_np_datetime_like
 from xarray.core.duck_array_ops import asarray
 from xarray.core.formatting import first_n_items, format_timestamp, last_item
-from xarray.core.parallelcompat import T_ChunkedArray, get_chunked_array_type
 from xarray.core.pdcompat import nanosecond_precision_timestamp
-from xarray.core.pycompat import is_chunked_array, is_duck_dask_array
 from xarray.core.utils import emit_user_level_warning
 from xarray.core.variable import Variable
+from xarray.namedarray.parallelcompat import T_ChunkedArray, get_chunked_array_type
+from xarray.namedarray.pycompat import is_chunked_array
+from xarray.namedarray.utils import is_duck_dask_array
 
 try:
     import cftime

xarray/coding/variables.py

@@ -11,9 +11,9 @@
 import pandas as pd
 
 from xarray.core import dtypes, duck_array_ops, indexing
-from xarray.core.parallelcompat import get_chunked_array_type
-from xarray.core.pycompat import is_chunked_array
 from xarray.core.variable import Variable
+from xarray.namedarray.parallelcompat import get_chunked_array_type
+from xarray.namedarray.pycompat import is_chunked_array
 
 if TYPE_CHECKING:
     T_VarTuple = tuple[tuple[Hashable, ...], Any, dict, dict]
@@ -163,7 +163,7 @@ def lazy_elemwise_func(array, func: Callable, dtype: np.typing.DTypeLike):
     if is_chunked_array(array):
         chunkmanager = get_chunked_array_type(array)
 
-        return chunkmanager.map_blocks(func, array, dtype=dtype)
+        return chunkmanager.map_blocks(func, array, dtype=dtype)  # type: ignore[arg-type]
     else:
         return _ElementwiseFunctionArray(array, func, dtype)
 

xarray/conventions.py

@@ -14,9 +14,9 @@
     _contains_datetime_like_objects,
     contains_cftime_datetimes,
 )
-from xarray.core.pycompat import is_duck_dask_array
 from xarray.core.utils import emit_user_level_warning
 from xarray.core.variable import IndexVariable, Variable
+from xarray.namedarray.utils import is_duck_dask_array
 
 CF_RELATED_DATA = (
     "bounds",

xarray/convert.py

@@ -10,7 +10,7 @@
 from xarray.core import duck_array_ops
 from xarray.core.dataarray import DataArray
 from xarray.core.dtypes import get_fill_value
-from xarray.core.pycompat import array_type
+from xarray.namedarray.pycompat import array_type
 
 iris_forbidden_keys = {
     "standard_name",

xarray/core/accessor_dt.py

@@ -13,9 +13,9 @@
     is_np_datetime_like,
     is_np_timedelta_like,
 )
-from xarray.core.pycompat import is_duck_dask_array
 from xarray.core.types import T_DataArray
 from xarray.core.variable import IndexVariable
+from xarray.namedarray.utils import is_duck_dask_array
 
 if TYPE_CHECKING:
     from numpy.typing import DTypeLike

xarray/core/arithmetic.py

@@ -15,12 +15,9 @@
     VariableOpsMixin,
 )
 from xarray.core.common import ImplementsArrayReduce, ImplementsDatasetReduce
-from xarray.core.ops import (
-    IncludeNumpySameMethods,
-    IncludeReduceMethods,
-)
+from xarray.core.ops import IncludeNumpySameMethods, IncludeReduceMethods
 from xarray.core.options import OPTIONS, _get_keep_attrs
-from xarray.core.pycompat import is_duck_array
+from xarray.namedarray.utils import is_duck_array
 
 
 class SupportsArithmetic:

xarray/core/common.py

@@ -13,15 +13,15 @@
 from xarray.core import dtypes, duck_array_ops, formatting, formatting_html, ops
 from xarray.core.indexing import BasicIndexer, ExplicitlyIndexed
 from xarray.core.options import OPTIONS, _get_keep_attrs
-from xarray.core.parallelcompat import get_chunked_array_type, guess_chunkmanager
-from xarray.core.pycompat import is_chunked_array
 from xarray.core.utils import (
     Frozen,
     either_dict_or_kwargs,
     emit_user_level_warning,
     is_scalar,
 )
 from xarray.namedarray.core import _raise_if_any_duplicate_dimensions
+from xarray.namedarray.parallelcompat import get_chunked_array_type, guess_chunkmanager
+from xarray.namedarray.pycompat import is_chunked_array
 
 try:
     import cftime

xarray/core/computation.py

@@ -22,11 +22,11 @@
 from xarray.core.indexes import Index, filter_indexes_from_coords
 from xarray.core.merge import merge_attrs, merge_coordinates_without_align
 from xarray.core.options import OPTIONS, _get_keep_attrs
-from xarray.core.parallelcompat import get_chunked_array_type
-from xarray.core.pycompat import is_chunked_array, is_duck_dask_array
 from xarray.core.types import Dims, T_DataArray
-from xarray.core.utils import is_dict_like, is_scalar, parse_dims
+from xarray.core.utils import is_dict_like, is_duck_dask_array, is_scalar, parse_dims
 from xarray.core.variable import Variable
+from xarray.namedarray.parallelcompat import get_chunked_array_type
+from xarray.namedarray.pycompat import is_chunked_array
 from xarray.util.deprecation_helpers import deprecate_dims
 
 if TYPE_CHECKING:

xarray/core/dataarray.py

@@ -64,14 +64,14 @@
     _default,
     either_dict_or_kwargs,
     hashable,
+    infix_dims,
 )
 from xarray.core.variable import (
     IndexVariable,
     Variable,
     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, deprecate_dims
@@ -85,7 +85,6 @@
     from xarray.backends import ZarrStore
     from xarray.backends.api import T_NetcdfEngine, T_NetcdfTypes
     from xarray.core.groupby import DataArrayGroupBy
-    from xarray.core.parallelcompat import ChunkManagerEntrypoint
     from xarray.core.resample import DataArrayResample
     from xarray.core.rolling import DataArrayCoarsen, DataArrayRolling
     from xarray.core.types import (
@@ -108,6 +107,7 @@
         T_Xarray,
     )
     from xarray.core.weighted import DataArrayWeighted
+    from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint
 
     T_XarrayOther = TypeVar("T_XarrayOther", bound=Union["DataArray", Dataset])
 

xarray/core/dataset.py

@@ -63,7 +63,6 @@
     assert_coordinate_consistent,
     create_coords_with_default_indexes,
 )
-from xarray.core.daskmanager import DaskManager
 from xarray.core.duck_array_ops import datetime_to_numeric
 from xarray.core.indexes import (
     Index,
@@ -86,13 +85,6 @@
 )
 from xarray.core.missing import get_clean_interp_index
 from xarray.core.options import OPTIONS, _get_keep_attrs
-from xarray.core.parallelcompat import get_chunked_array_type, guess_chunkmanager
-from xarray.core.pycompat import (
-    array_type,
-    is_chunked_array,
-    is_duck_array,
-    is_duck_dask_array,
-)
 from xarray.core.types import (
     QuantileMethods,
     Self,
@@ -114,6 +106,10 @@
     drop_dims_from_indexers,
     either_dict_or_kwargs,
     emit_user_level_warning,
+    infix_dims,
+    is_dict_like,
+    is_duck_array,
+    is_duck_dask_array,
     is_scalar,
     maybe_wrap_array,
 )
@@ -124,7 +120,8 @@
     broadcast_variables,
     calculate_dimensions,
 )
-from xarray.namedarray.utils import infix_dims, is_dict_like
+from xarray.namedarray.parallelcompat import get_chunked_array_type, guess_chunkmanager
+from xarray.namedarray.pycompat import array_type, is_chunked_array
 from xarray.plot.accessor import DatasetPlotAccessor
 from xarray.util.deprecation_helpers import _deprecate_positional_args
 
@@ -138,7 +135,6 @@
     from xarray.core.dataarray import DataArray
     from xarray.core.groupby import DatasetGroupBy
     from xarray.core.merge import CoercibleMapping, CoercibleValue, _MergeResult
-    from xarray.core.parallelcompat import ChunkManagerEntrypoint
     from xarray.core.resample import DatasetResample
     from xarray.core.rolling import DatasetCoarsen, DatasetRolling
     from xarray.core.types import (
@@ -164,6 +160,7 @@
         T_Xarray,
     )
     from xarray.core.weighted import DatasetWeighted
+    from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint
 
 
 # list of attributes of pd.DatetimeIndex that are ndarrays of time info
@@ -292,6 +289,9 @@ def _maybe_chunk(
     chunked_array_type: str | ChunkManagerEntrypoint | None = None,
     from_array_kwargs=None,
 ):
+
+    from xarray.namedarray.daskmanager import DaskManager
+
     if chunks is not None:
         chunks = {dim: chunks[dim] for dim in var.dims if dim in chunks}
 
@@ -842,7 +842,9 @@ def load(self, **kwargs) -> Self:
             chunkmanager = get_chunked_array_type(*lazy_data.values())
 
             # evaluate all the chunked arrays simultaneously
-            evaluated_data = chunkmanager.compute(*lazy_data.values(), **kwargs)
+            evaluated_data: tuple[np.ndarray[Any, Any], ...] = chunkmanager.compute(
+                *lazy_data.values(), **kwargs
+            )
 
             for k, data in zip(lazy_data, evaluated_data):
                 self.variables[k].data = data

xarray/core/duck_array_ops.py

@@ -33,11 +33,12 @@
 from numpy.lib.stride_tricks import sliding_window_view  # noqa
 from packaging.version import Version
 
-from xarray.core import dask_array_ops, dtypes, nputils, pycompat
+from xarray.core import dask_array_ops, dtypes, nputils
 from xarray.core.options import OPTIONS
-from xarray.core.parallelcompat import get_chunked_array_type, is_chunked_array
-from xarray.core.pycompat import array_type, is_duck_dask_array
-from xarray.core.utils import is_duck_array, module_available
+from xarray.core.utils import is_duck_array, is_duck_dask_array, module_available
+from xarray.namedarray import pycompat
+from xarray.namedarray.parallelcompat import get_chunked_array_type
+from xarray.namedarray.pycompat import array_type, is_chunked_array
 
 # remove once numpy 2.0 is the oldest supported version
 if module_available("numpy", minversion="2.0.0.dev0"):

xarray/core/formatting.py

@@ -20,8 +20,8 @@
 from xarray.core.duck_array_ops import array_equiv, astype
 from xarray.core.indexing import MemoryCachedArray
 from xarray.core.options import OPTIONS, _get_boolean_with_default
-from xarray.core.pycompat import array_type, to_duck_array, to_numpy
 from xarray.core.utils import is_duck_array
+from xarray.namedarray.pycompat import array_type, to_duck_array, to_numpy
 
 if TYPE_CHECKING:
     from xarray.core.coordinates import AbstractCoordinates