import normalize_axis_index from numpy.lib on numpy>=2 (#364)

* import `normalize_axis_index` from `numpy.lib` on `numpy>=2`

* import the right thing

flox/xrutils.py

@@ -8,7 +8,6 @@
 
 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
 
 try:
@@ -25,6 +24,37 @@
  dask_array_type = () # type: ignore[assignment, misc]
 
 
+def module_available(module: str, minversion: Optional[str] = None) -> bool:
+ """Checks whether a module is installed without importing it.
+
+ Use this for a lightweight check and lazy imports.
+
+ Parameters
+ ----------
+ module : str
+ Name of the module.
+
+ Returns
+ -------
+ available : bool
+ Whether the module is installed.
+ """
+ has = importlib.util.find_spec(module) is not None
+ if has:
+ mod = importlib.import_module(module)
+ return Version(mod.__version__) >= Version(minversion) if minversion is not None else True
+ else:
+ return False
+
+
+if module_available("numpy", minversion="2.0.0"):
+ from numpy.lib.array_utils import ( # type: ignore[import-not-found]
+ normalize_axis_index,
+ )
+else:
+ from numpy.core.numeric import normalize_axis_index # type: ignore[attr-defined]
+
+
 def asarray(data, xp=np):
  return data if is_duck_array(data) else xp.asarray(data)
 
@@ -349,26 +379,3 @@ def nanlast(values, axis, keepdims=False):
  return np.expand_dims(result, axis=axis)
  else:
  return result
-
-
-def module_available(module: str, minversion: Optional[str] = None) -> bool:
- """Checks whether a module is installed without importing it.
-
- Use this for a lightweight check and lazy imports.
-
- Parameters
- ----------
- module : str
- Name of the module.
-
- Returns
- -------
- available : bool
- Whether the module is installed.
- """
- has = importlib.util.find_spec(module) is not None
- if has:
- mod = importlib.import_module(module)
- return Version(mod.__version__) >= Version(minversion) if minversion is not None else True
- else:
- return False