Fix upcasting with python builtin numbers and numpy 2

xarray/core/dtypes.py

@@ -215,31 +215,6 @@ def isdtype(dtype, kind: str | tuple[str, ...], xp=None) -> bool:
         return xp.isdtype(dtype, kind)
 
 
-def is_scalar_type(t):
-    return isinstance(t, (bool, int, float, complex, str, bytes))
-
-
-def _future_array_api_result_type(*arrays_and_dtypes, xp):
-    strongly_dtyped = [t for t in arrays_and_dtypes if not is_scalar_type(t)]
-    weakly_dtyped = [t for t in arrays_and_dtypes if is_scalar_type(t)]
-
-    dtype = xp.result_type(*strongly_dtyped)
-    if not weakly_dtyped:
-        return dtype
-
-    possible_dtypes = {
-        complex: "complex64",
-        float: "float32",
-        int: "int8",
-        bool: "bool",
-        str: "str",
-        bytes: "bytes",
-    }
-    dtypes = [possible_dtypes.get(type(x), "object") for x in weakly_dtyped]
-
-    return xp.result_type(dtype, *dtypes)
-
-
 def preprocess_scalar_types(t):
     if isinstance(t, (str, bytes)):
         return type(t)
@@ -266,12 +241,17 @@ def result_type(
     -------
     numpy.dtype for the result.
     """
+    # TODO (keewis): replace `npcompat.result_type` with `xp.result_type` once we can
+    # require a version of the Array API that supports passing scalars to it.
     from xarray.core.duck_array_ops import get_array_namespace
 
     if xp is None:
         xp = get_array_namespace(arrays_and_dtypes)
 
-    types = {xp.result_type(preprocess_scalar_types(t)) for t in arrays_and_dtypes}
+    types = {
+        npcompat.result_type(preprocess_scalar_types(t), xp=xp)
+        for t in arrays_and_dtypes
+    }
     if any(isinstance(t, np.dtype) for t in types):
         # only check if there's numpy dtypes – the array API does not
         # define the types we're checking for
@@ -281,9 +261,4 @@ def result_type(
             ):
                 return np.dtype(object)
 
-    if xp is np or any(
-        isinstance(getattr(t, "dtype", t), np.dtype) for t in arrays_and_dtypes
-    ):
-        return xp.result_type(*map(preprocess_scalar_types, arrays_and_dtypes))
-    else:
-        return _future_array_api_result_type(*arrays_and_dtypes, xp=xp)
+    return npcompat.result_type(*map(preprocess_scalar_types, arrays_and_dtypes), xp=xp)

xarray/core/duck_array_ops.py

@@ -251,7 +251,7 @@ def asarray(data, xp=np, dtype=None):
         return xp.astype(converted, dtype)
 
 
-def as_shared_dtype(scalars_or_arrays, xp=np):
+def as_shared_dtype(scalars_or_arrays, xp=None):
     """Cast a arrays to a shared dtype using xarray's type promotion rules."""
     if any(is_extension_array_dtype(x) for x in scalars_or_arrays):
         extension_array_types = [
@@ -272,6 +272,8 @@ def as_shared_dtype(scalars_or_arrays, xp=np):
         import cupy as cp
 
         xp = cp
+    elif xp is None:
+        xp = get_array_namespace(scalars_or_arrays)
 
     # Pass arrays directly instead of dtypes to result_type so scalars
     # get handled properly.

xarray/core/npcompat.py

@@ -29,11 +29,12 @@
 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import numpy as np
+
 try:
     # requires numpy>=2.0
     from numpy import isdtype  # type: ignore[attr-defined,unused-ignore]
 except ImportError:
-    import numpy as np
 
     dtype_kinds = {
         "bool": np.bool_,
@@ -58,3 +59,43 @@ def isdtype(dtype, kind):
             return any(isinstance(dtype, kind) for kind in translated_kinds)
         else:
             return any(np.issubdtype(dtype, kind) for kind in translated_kinds)
+
+
+def is_weak_scalar_type(t):
+    return isinstance(t, (bool, int, float, complex, str, bytes))
+
+
+def _future_array_api_result_type(*arrays_and_dtypes, xp):
+    strongly_dtyped = [t for t in arrays_and_dtypes if not is_weak_scalar_type(t)]
+    weakly_dtyped = [t for t in arrays_and_dtypes if is_weak_scalar_type(t)]
+
+    if not strongly_dtyped:
+        strongly_dtyped = [
+            xp.asarray(x) if not isinstance(x, type) else x for x in weakly_dtyped
+        ]
+        weakly_dtyped = []
+
+    dtype = xp.result_type(*strongly_dtyped)
+    if not weakly_dtyped:
+        return dtype
+
+    possible_dtypes = {
+        complex: "complex64",
+        float: "float32",
+        int: "int8",
+        bool: "bool",
+        str: "str",
+        bytes: "bytes",
+    }
+    dtypes = [possible_dtypes.get(type(x), "object") for x in weakly_dtyped]
+
+    return xp.result_type(dtype, *dtypes)
+
+
+def result_type(*arrays_and_dtypes, xp) -> np.dtype:
+    if xp is np or any(
+        isinstance(getattr(t, "dtype", t), np.dtype) for t in arrays_and_dtypes
+    ):
+        return xp.result_type(*arrays_and_dtypes)
+    else:
+        return _future_array_api_result_type(*arrays_and_dtypes, xp=xp)

xarray/tests/test_array_api.py

@@ -137,9 +137,6 @@ def test_unstack(arrays: tuple[xr.DataArray, xr.DataArray]) -> None:
     assert_equal(actual, expected)
 
 
-@pytest.mark.xfail(
-    reason="We don't know how to handle python scalars with array API arrays. This is strictly prohibited by the array API"
-)
 def test_where() -> None:
     np_arr = xr.DataArray(np.array([1, 0]), dims="x")
     xp_arr = xr.DataArray(xp.asarray([1, 0]), dims="x")