import from the new location of normalize_axis_index if possible (#8483)

* import from the new location of `normalize_axis_index` if possible

* disable build isolation for `cftime` [skip-ci] [skip-rtd]

* actually apply the `--no-build-isolation` [skip-ci] [skip-rtd]

* install `cython`, which is needed to build `cftime`

* use `micromamba` instead of `conda` for removing packages

* switch the main uninstall back to `conda`

* temporarily remove `cf_units`, `h5py`, and `netcdf4`

* also remove `hdf5` [skip-ci]

* try building `numcodecs` from github and without build isolation

* build `cftime` separately from `numcodecs` and install more deps [skip-ci]

* don't uninstall packages that are already removed [skip-ci]

* also build `bottleneck` without build isolation

* use `module_available` instead of eagerly importing

* add a `minversion` kwarg to `module_available`

* use `module_available` in `nputils` as well

* more type ignores

* move the type ignores to the import

* ignore complaints about unused ignores

(the point here is compatibility between multiple versions)

* a couple of other `type: ignore` comments (only relevant for nightly)

Author: keewis

ci/install-upstream-wheels.sh

@@ -1,26 +1,25 @@
 #!/usr/bin/env bash
 
+# install cython for building cftime without build isolation
+micromamba install "cython>=0.29.20" py-cpuinfo
 # temporarily (?) remove numbagg and numba
-pip uninstall -y numbagg
-conda uninstall -y numba
+micromamba remove -y numba numbagg
+# temporarily remove backends
+micromamba remove -y cf_units h5py hdf5 netcdf4
 # forcibly remove packages to avoid artifacts
 conda uninstall -y --force \
     numpy \
     scipy \
     pandas \
-    matplotlib \
-    dask \
     distributed \
     fsspec \
     zarr \
     cftime \
     packaging \
     pint \
     bottleneck \
-    sparse \
     flox \
-    h5netcdf \
-    xarray
+    numcodecs
 # to limit the runtime of Upstream CI
 python -m pip install \
     -i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple \
@@ -31,19 +30,34 @@ python -m pip install \
     scipy \
     matplotlib \
     pandas
+# without build isolation for packages compiling against numpy
+# TODO: remove once there are `numpy>=2.0` builds for numcodecs and cftime
+python -m pip install \
+    --no-deps \
+    --upgrade \
+    --no-build-isolation \
+    git+https://github.com/Unidata/cftime
+python -m pip install \
+    --no-deps \
+    --upgrade \
+    --no-build-isolation \
+    git+https://github.com/zarr-developers/numcodecs
+python -m pip install \
+    --no-deps \
+    --upgrade \
+    --no-build-isolation \
+    git+https://github.com/pydata/bottleneck
 python -m pip install \
     --no-deps \
     --upgrade \
     git+https://github.com/dask/dask \
     git+https://github.com/dask/distributed \
     git+https://github.com/zarr-developers/zarr \
-    git+https://github.com/Unidata/cftime \
     git+https://github.com/pypa/packaging \
     git+https://github.com/hgrecco/pint \
-    git+https://github.com/pydata/bottleneck \
     git+https://github.com/pydata/sparse \
     git+https://github.com/intake/filesystem_spec \
     git+https://github.com/SciTools/nc-time-axis \
     git+https://github.com/xarray-contrib/flox \
-    git+https://github.com/h5netcdf/h5netcdf \
     git+https://github.com/dgasmith/opt_einsum
+    # git+https://github.com/h5netcdf/h5netcdf

xarray/core/duck_array_ops.py

@@ -29,7 +29,6 @@
     zeros_like,  # noqa
 )
 from numpy import concatenate as _concatenate
-from numpy.core.multiarray import normalize_axis_index  # type: ignore[attr-defined]
 from numpy.lib.stride_tricks import sliding_window_view  # noqa
 from packaging.version import Version
 
@@ -39,6 +38,17 @@
 from xarray.core.pycompat import array_type, is_duck_dask_array
 from xarray.core.utils import is_duck_array, module_available
 
+# remove once numpy 2.0 is the oldest supported version
+if module_available("numpy", minversion="2.0.0.dev0"):
+    from numpy.lib.array_utils import (  # type: ignore[import-not-found,unused-ignore]
+        normalize_axis_index,
+    )
+else:
+    from numpy.core.multiarray import (  # type: ignore[attr-defined,no-redef,unused-ignore]
+        normalize_axis_index,
+    )
+
+
 dask_available = module_available("dask")
 
 

xarray/core/nputils.py

@@ -5,17 +5,26 @@
 
 import numpy as np
 import pandas as pd
-from numpy.core.multiarray import normalize_axis_index  # type: ignore[attr-defined]
 from packaging.version import Version
 
 from xarray.core import pycompat
 from xarray.core.utils import module_available
 
+# remove once numpy 2.0 is the oldest supported version
+if module_available("numpy", minversion="2.0.0.dev0"):
+    from numpy.lib.array_utils import (  # type: ignore[import-not-found,unused-ignore]
+        normalize_axis_index,
+    )
+else:
+    from numpy.core.multiarray import (  # type: ignore[attr-defined,no-redef,unused-ignore]
+        normalize_axis_index,
+    )
+
 # remove once numpy 2.0 is the oldest supported version
 try:
     from numpy.exceptions import RankWarning  # type: ignore[attr-defined,unused-ignore]
 except ImportError:
-    from numpy import RankWarning
+    from numpy import RankWarning  # type: ignore[attr-defined,no-redef,unused-ignore]
 
 from xarray.core.options import OPTIONS
 from xarray.core.pycompat import is_duck_array

xarray/core/utils.py

@@ -73,6 +73,7 @@
 
 import numpy as np
 import pandas as pd
+from packaging.version import Version
 
 if TYPE_CHECKING:
     from xarray.core.types import Dims, ErrorOptionsWithWarn, T_DuckArray
@@ -1194,7 +1195,7 @@ def contains_only_chunked_or_numpy(obj) -> bool:
     )
 
 
-def module_available(module: str) -> bool:
+def module_available(module: str, minversion: str | None = None) -> bool:
     """Checks whether a module is installed without importing it.
 
     Use this for a lightweight check and lazy imports.
@@ -1203,13 +1204,23 @@ def module_available(module: str) -> bool:
     ----------
     module : str
         Name of the module.
+    minversion : str, optional
+        Minimum version of the module
 
     Returns
     -------
     available : bool
         Whether the module is installed.
     """
-    return importlib.util.find_spec(module) is not None
+    if importlib.util.find_spec(module) is None:
+        return False
+
+    if minversion is not None:
+        version = importlib.metadata.version(module)
+
+        return Version(version) >= Version(minversion)
+
+    return True
 
 
 def find_stack_level(test_mode=False) -> int: