implement additional request checks

Author: toddrjen

xarray/core/computation.py

@@ -1322,7 +1322,7 @@ def _calc_idxminmax(
     """Apply common operations for idxmin and idxmax."""
     # This function doesn't make sense for scalars so don't try
     if not array.ndim:
-        ValueError("This function does not apply for scalars")
+        raise ValueError("This function does not apply for scalars")
 
     if dim is not None:
         pass  # Use the dim if available
@@ -1333,13 +1333,13 @@ def _calc_idxminmax(
         # The dim is not specified and ambiguous.  Don't guess.
         raise ValueError("Must supply 'dim' argument for multidimensional arrays")
 
-    if dim in array.coords:
-        pass  # This is okay
-    else:
+    if dim not in array.dims:
+        raise KeyError(f'Dimension "{dim}" not in dimension')
+    if dim not in array.coords:
         raise KeyError(f'Dimension "{dim}" does not have coordinates')
 
     # These are dtypes with NaN values argmin and argmax can handle
-    na_dtypes = "cf0"
+    na_dtypes = "cfO"
 
     if skipna or (skipna is None and array.dtype.kind in na_dtypes):
         # Need to skip NaN values since argmin and argmax can't handle them

xarray/core/dataset.py

@@ -6,6 +6,7 @@
 from collections import defaultdict
 from html import escape
 from numbers import Number
+from operator import methodcaller
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
@@ -6003,12 +6004,14 @@ def idxmin(
         DataArray.idxmin, Dataset.idxmax, Dataset.min, Dataset.argmin
         """
         return self.map(
-            "idxmin",
-            dim=dim,
-            skipna=skipna,
-            promote=promote,
-            keep_attrs=keep_attrs,
-            **kwargs,
+            methodcaller(
+                "idxmin",
+                dim=dim,
+                skipna=skipna,
+                promote=promote,
+                keep_attrs=keep_attrs,
+                **kwargs,
+            ),
         )
 
     def idxmax(
@@ -6062,12 +6065,14 @@ def idxmax(
         DataArray.idxmax, Dataset.idxmin, Dataset.max, Dataset.argmax
         """
         return self.map(
-            "idxmax",
-            dim=dim,
-            skipna=skipna,
-            promote=promote,
-            keep_attrs=keep_attrs,
-            **kwargs,
+            methodcaller(
+                "idxmax",
+                dim=dim,
+                skipna=skipna,
+                promote=promote,
+                keep_attrs=keep_attrs,
+                **kwargs,
+            ),
         )
 
 

xarray/tests/test_dataarray.py

@@ -4298,6 +4298,12 @@ def setup(self):
         (np.array([0, 1, 2, 0, -2, -4, 2]), 5, 2, None),
         (np.array([0.0, 1.0, 2.0, 0.0, -2.0, -4.0, 2.0]), 5, 2, None),
         (np.array([1.0, np.NaN, 2.0, np.NaN, -2.0, -4.0, 2.0]), 5, 2, 1),
+        (
+            np.array([1.0, np.NaN, 2.0, np.NaN, -2.0, -4.0, 2.0]).astype("object"),
+            5,
+            2,
+            1,
+        ),
         (np.array([np.NaN, np.NaN]), np.NaN, np.NaN, 0),
         (
             np.array(
@@ -4329,7 +4335,7 @@ def test_min(self, x, minindex, maxindex, nanindex):
         assert_identical(result1, expected1)
 
         result2 = ar.min(skipna=False)
-        if nanindex is not None:
+        if nanindex is not None and ar.dtype.kind != "O":
             expected2 = ar.isel(x=nanindex, drop=True)
             expected2.attrs = {}
         else:
@@ -4355,7 +4361,7 @@ def test_max(self, x, minindex, maxindex, nanindex):
         assert_identical(result1, expected1)
 
         result2 = ar.max(skipna=False)
-        if nanindex is not None:
+        if nanindex is not None and ar.dtype.kind != "O":
             expected2 = ar.isel(x=nanindex, drop=True)
             expected2.attrs = {}
         else:
@@ -4384,7 +4390,7 @@ def test_argmin(self, x, minindex, maxindex, nanindex):
         assert_identical(result1, expected1)
 
         result2 = ar.argmin(skipna=False)
-        if nanindex is not None:
+        if nanindex is not None and ar.dtype.kind != "O":
             expected2 = indarr.isel(x=nanindex, drop=True)
             expected2.attrs = {}
         else:
@@ -4413,7 +4419,7 @@ def test_argmax(self, x, minindex, maxindex, nanindex):
         assert_identical(result1, expected1)
 
         result2 = ar.argmax(skipna=False)
-        if nanindex is not None:
+        if nanindex is not None and ar.dtype.kind != "O":
             expected2 = indarr.isel(x=nanindex, drop=True)
             expected2.attrs = {}
         else:
@@ -4448,7 +4454,7 @@ def test_idxmin(self, x, minindex, maxindex, nanindex):
         assert_identical(result2, expected1)
 
         result3 = ar.idxmin(skipna=False)
-        if nanindex is not None:
+        if nanindex is not None and ar.dtype.kind != "O":
             expected2 = coordarr.isel(x=nanindex, drop=True).astype("float")
             expected2.name = "x"
             expected2.attrs = {}
@@ -4461,7 +4467,7 @@ def test_idxmin(self, x, minindex, maxindex, nanindex):
             with pytest.raises(TypeError):
                 ar.idxmin(promote=False)
 
-        if nanindex is None:
+        if nanindex is None or ar.dtype.kind == "O":
             expected3 = coordarr.isel(x=minindex, drop=True)
             expected3.name = "x"
 
@@ -4480,6 +4486,13 @@ def test_idxmin(self, x, minindex, maxindex, nanindex):
             result7 = ar.idxmin(skipna=False, promote=None)
             assert_identical(result7, expected4)
 
+        with pytest.raises(KeyError):
+            ar.idxmin(dim="spam")
+
+        ar2 = xr.DataArray(5)
+        with pytest.raises(ValueError):
+            ar2.idxmin()
+
     def test_idxmax(self, x, minindex, maxindex, nanindex):
         ar = xr.DataArray(
             x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs,
@@ -4507,7 +4520,7 @@ def test_idxmax(self, x, minindex, maxindex, nanindex):
         assert_identical(result2, expected1)
 
         result3 = ar.idxmax(skipna=False)
-        if nanindex is not None:
+        if nanindex is not None and ar.dtype.kind != "O":
             expected2 = coordarr.isel(x=nanindex, drop=True).astype("float")
             expected2.name = "x"
             expected2.attrs = {}
@@ -4529,7 +4542,7 @@ def test_idxmax(self, x, minindex, maxindex, nanindex):
             result5 = ar2.idxmax(promote=None)
             assert_identical(result5, expected1)
 
-        if nanindex is None:
+        if nanindex is None or ar.dtype.kind == "O":
             expected3 = coordarr.isel(x=maxindex, drop=True)
             expected3.name = "x"
 
@@ -4548,6 +4561,13 @@ def test_idxmax(self, x, minindex, maxindex, nanindex):
             result9 = ar.idxmax(skipna=False, promote=None)
             assert_identical(result9, expected4)
 
+        with pytest.raises(KeyError):
+            ar.idxmin(dim="spam")
+
+        ar2 = xr.DataArray(5)
+        with pytest.raises(ValueError):
+            ar2.idxmin()
+
 
 @pytest.mark.parametrize(
     "x, minindex, maxindex, nanindex",
@@ -4576,6 +4596,18 @@ def test_idxmax(self, x, minindex, maxindex, nanindex):
             [0, 2, np.NaN],
             [None, 1, 0],
         ),
+        (
+            np.array(
+                [
+                    [2.0, 1.0, 2.0, 0.0, -2.0, -4.0, 2.0],
+                    [-4.0, np.NaN, 2.0, np.NaN, -2.0, -4.0, 2.0],
+                    [np.NaN] * 7,
+                ]
+            ).astype("object"),
+            [5, 0, np.NaN],
+            [0, 2, np.NaN],
+            [None, 1, 0],
+        ),
         (
             np.array(
                 [
@@ -4617,7 +4649,10 @@ def test_min(self, x, minindex, maxindex, nanindex):
         result2 = ar.min(axis=1)
         assert_identical(result2, expected1)
 
-        minindex = [x if y is None else y for x, y in zip(minindex, nanindex)]
+        minindex = [
+            x if y is None or ar.dtype.kind == "O" else y
+            for x, y in zip(minindex, nanindex)
+        ]
         expected2 = [
             ar.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex)
         ]
@@ -4653,7 +4688,10 @@ def test_max(self, x, minindex, maxindex, nanindex):
         result2 = ar.max(axis=1)
         assert_identical(result2, expected1)
 
-        maxindex = [x if y is None else y for x, y in zip(maxindex, nanindex)]
+        maxindex = [
+            x if y is None or ar.dtype.kind == "O" else y
+            for x, y in zip(maxindex, nanindex)
+        ]
         expected2 = [
             ar.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex)
         ]
@@ -4696,7 +4734,10 @@ def test_argmin(self, x, minindex, maxindex, nanindex):
         expected1.attrs = self.attrs
         assert_identical(result2, expected1)
 
-        minindex = [x if y is None else y for x, y in zip(minindex, nanindex)]
+        minindex = [
+            x if y is None or ar.dtype.kind == "O" else y
+            for x, y in zip(minindex, nanindex)
+        ]
         expected2 = [
             indarr.isel(y=yi).isel(x=indi, drop=True)
             for yi, indi in enumerate(minindex)
@@ -4740,7 +4781,10 @@ def test_argmax(self, x, minindex, maxindex, nanindex):
         expected1.attrs = self.attrs
         assert_identical(result2, expected1)
 
-        maxindex = [x if y is None else y for x, y in zip(maxindex, nanindex)]
+        maxindex = [
+            x if y is None or ar.dtype.kind == "O" else y
+            for x, y in zip(maxindex, nanindex)
+        ]
         expected2 = [
             indarr.isel(y=yi).isel(x=indi, drop=True)
             for yi, indi in enumerate(maxindex)
@@ -4782,7 +4826,10 @@ def test_idxmin(self, x, minindex, maxindex, nanindex):
         expected1.attrs = self.attrs
         assert_identical(result1, expected1)
 
-        minindex = [x if y is None else y for x, y in zip(minindex, nanindex)]
+        minindex = [
+            x if y is None or ar.dtype.kind == "O" else y
+            for x, y in zip(minindex, nanindex)
+        ]
         expected2 = [
             coordarr.isel(y=yi).isel(x=indi, drop=True)
             for yi, indi in enumerate(minindex)
@@ -4835,6 +4882,9 @@ def test_idxmin(self, x, minindex, maxindex, nanindex):
         with pytest.raises(KeyError):
             ar.idxmin(dim="y")
 
+        with pytest.raises(KeyError):
+            ar.idxmax(dim="spam")
+
     def test_idxmax(self, x, minindex, maxindex, nanindex):
         ar = xr.DataArray(
             x,
@@ -4865,7 +4915,10 @@ def test_idxmax(self, x, minindex, maxindex, nanindex):
         expected1.attrs = self.attrs
         assert_identical(result1, expected1)
 
-        maxindex = [x if y is None else y for x, y in zip(maxindex, nanindex)]
+        maxindex = [
+            x if y is None or ar.dtype.kind == "O" else y
+            for x, y in zip(maxindex, nanindex)
+        ]
         expected2 = [
             coordarr.isel(y=yi).isel(x=indi, drop=True)
             for yi, indi in enumerate(maxindex)
@@ -4918,6 +4971,9 @@ def test_idxmax(self, x, minindex, maxindex, nanindex):
         with pytest.raises(KeyError):
             ar.idxmax(dim="y")
 
+        with pytest.raises(KeyError):
+            ar.idxmax(dim="spam")
+
 
 @pytest.fixture(params=[1])
 def da(request):