FEAT: made Array[] error message more useful (closes #875)

gdementen · gdementen · commit e9d8e73c959b · 2023-02-14T21:44:31.000+01:00
We now have some axes info in the error messages

also factorized
* parts of _guess_axis and _translate_axis_key_chunk
* _display_name, _axis_summary and _axes_summary methods in AxisCollection
and this fixed two corner cases for ambiguous labels on anonymous axes
diff --git a/doc/source/changes/version_0_34.rst.inc b/doc/source/changes/version_0_34.rst.inc
@@ -44,6 +44,14 @@ Miscellaneous improvements
 * added type hints for all remaining functions and methods which improves autocompletion in editors (such as PyCharm).
   Closes :issue:`864`.
 
+* made several error message more useful when trying to get an invalid subset of an array (closes :issue:`875`).
+  - when a key is not valid on any axis, the error message includes the array axes
+  - when a key is not valid for the axis specified by the user, the error message includes that axis labels
+  - when a label is ambiguous (valid on several axes), the error message includes the axes labels in addition to the
+    axes names
+  - when several parts of a key seem to target the same axis, the error message includes the bad key in
+    addition to the axis.
+
 * made :py:obj:`ipfp()` faster (the smaller the array, the larger the improvement)
   For example, for small arrays it is several times faster than before, for 1000x1000 arrays it is about 30% faster.
 
diff --git a/larray/core/axis.py b/larray/core/axis.py
@@ -2439,37 +2439,6 @@ def replace(self, axes_to_replace=None, new_axis=None, inplace=False, **kwargs)
         else:
             return AxisCollection(axes)
 
-    def _guess_axis(self, axis_key):
-        if isinstance(axis_key, Group):
-            group_axis = axis_key.axis
-            if group_axis is not None:
-                # we have axis information but not necessarily an Axis object from self.axes
-                real_axis = self[group_axis]
-                if group_axis is not real_axis:
-                    axis_key = axis_key.with_axis(real_axis)
-                return axis_key
-
-        # TODO: instead of checking all axes, we should have a big mapping
-        # (in AxisCollection or Array):
-        # label -> (axis, index)
-        # or possibly (for ambiguous labels)
-        # label -> {axis: index}
-        # but for Pandas, this wouldn't work, we'd need label -> axis
-        valid_axes = []
-        for axis in self:
-            try:
-                axis.index(axis_key)
-                valid_axes.append(axis)
-            except KeyError:
-                continue
-        if not valid_axes:
-            raise ValueError(f"{axis_key} is not a valid label for any axis\n{self.info}")
-        elif len(valid_axes) > 1:
-            valid_axes = ', '.join(a.name if a.name is not None else f'{{{self.axes.index(a)}}}'
-                                   for a in valid_axes)
-            raise ValueError(f'{axis_key} is ambiguous (valid in {valid_axes})')
-        return valid_axes[0][axis_key]
-
     def set_labels(self, axis=None, labels=None, inplace=False, **kwargs) -> 'AxisCollection':
         r"""Replaces the labels of one or several axes.
 
@@ -2676,9 +2645,44 @@ def index_first_compatible(axis):
         # -1 in to_remove are not a problem since enumerate starts at 0
         return AxisCollection([axis for i, axis in enumerate(self) if i not in to_remove])
 
+    def _translate_nice_key(self, axis_key):
+        # TODO: instead of checking all axes, we should have a big mapping (in AxisCollection):
+        #       label -> (axis, index) but for sparse/multi-index, this would not work, we'd need label -> axis
+        valid_axes = []
+        # TODO: use axis_key dtype to only check compatible axes
+        for axis in self:
+            try:
+                axis_pos_key = axis.index(axis_key)
+                valid_axes.append(axis)
+            except KeyError:
+                continue
+        if not valid_axes:
+            raise ValueError(f"{axis_key!r} is not a valid label for any axis:\n{self._axes_summary()}")
+        elif len(valid_axes) > 1:
+            raise ValueError(f'{axis_key!r} is ambiguous, it is valid in the following axes:\n'
+                             f'{self._axes_summary(valid_axes)}')
+        real_axis = valid_axes[0]
+        return real_axis, axis_pos_key
+
+    def _guess_axis(self, axis_key):
+        """
+        Translates any *single axis* key to an LGroup on the real axis.
+        """
+        if isinstance(axis_key, Group):
+            group_axis = axis_key.axis
+            if group_axis is not None:
+                # we have axis information but not necessarily an Axis object from self
+                real_axis = self[group_axis]
+                if group_axis is not real_axis:
+                    axis_key = axis_key.with_axis(real_axis)
+                return axis_key
+
+        real_axis, axis_pos_key = self._translate_nice_key(axis_key)
+        return real_axis[axis_key]
+
     def _translate_axis_key_chunk(self, axis_key):
         """
-        Translates *single axis* label-based key to an IGroup
+        Translates any *single axis* label-based key to an (axis, indices) pair.
 
         Parameters
         ----------
@@ -2717,7 +2721,8 @@ def _translate_axis_key_chunk(self, axis_key):
                 try:
                     axis_pos_key = real_axis.index(axis_key)
                 except KeyError:
-                    raise ValueError(f"{axis_key!r} is not a valid label for any axis")
+                    raise ValueError(f"{axis_key!r} is not a valid label for the {real_axis.name!r} axis "
+                                     f"with labels: {', '.join(repr(label) for label in real_axis.labels)}")
                 return real_axis, axis_pos_key
             except KeyError:
                 # axis associated with axis_key may not belong to self.
@@ -2726,26 +2731,7 @@ def _translate_axis_key_chunk(self, axis_key):
                 axis_key = axis_key.to_label()
 
         # otherwise we need to guess the axis
-        # TODO: instead of checking all axes, we should have a big mapping (in AxisCollection):
-        #       label -> (axis, index) but for sparse/multi-index, this would not work, we'd need label -> axis
-        valid_axes = []
-        # TODO: use axis_key dtype to only check compatible axes
-        for axis in self:
-            try:
-                axis_pos_key = axis.index(axis_key)
-                valid_axes.append(axis)
-            except KeyError:
-                continue
-        if not valid_axes:
-            raise ValueError(f"{axis_key!r} is not a valid label for any axis")
-        elif len(valid_axes) > 1:
-            # TODO: make an AxisCollection.display_name(axis) method out of this
-            # valid_axes = ', '.join(self.display_name(axis) for a in valid_axes)
-            valid_axes = ', '.join(a.name if a.name is not None else f'{{{self.index(a)}}}'
-                                   for a in valid_axes)
-            raise ValueError(f'{axis_key} is ambiguous (valid in {valid_axes})')
-        real_axis = valid_axes[0]
-        return real_axis, axis_pos_key
+        return self._translate_nice_key(axis_key)
 
     def _translate_axis_key(self, axis_key):
         """
@@ -2878,7 +2864,7 @@ def _key_to_axis_indices_dict(self, key):
         if has_duplicates(axis for axis, axis_key in key_items):
             dupe_axes = duplicates(axis for axis, axis_key in key_items)
             dupe_axes_str = ', '.join(str(axis) for axis in dupe_axes)
-            raise ValueError(f"key has several values for axis: {dupe_axes_str}\n{key}")
+            raise ValueError(f"key has several values for axis: {dupe_axes_str}\nkey: {key}")
 
         # ((axis, indices), (axis, indices), ...) -> dict
         return dict(key_items)
@@ -3093,6 +3079,27 @@ def names(self) -> List[str]:
         """
         return [axis.name for axis in self._list]
 
+    # providing idx is just an optimization to avoid the relatively expensive self.index(axis)
+    def _display_name(self, axis, idx=None):
+        if axis.name is None:
+            if idx is None:
+                idx = self.index(axis)
+            name = f'{{{idx}}}'
+        else:
+            # using str() because name can be an integer
+            name = str(axis.name)
+        return (name + '*') if axis.iswildcard else name
+
+    def _axes_summary(self, axes=None):
+        def axis_summary(axis, idx=None):
+            return f" {self._display_name(axis, idx)} [{len(axis)}]: {axis.labels_summary()}"
+
+        if axes is None:
+            parts = [axis_summary(axis, i) for i, axis in enumerate(self._list)]
+        else:
+            parts = [axis_summary(axis) for axis in axes]
+        return '\n'.join(parts)
+
     @property
     def display_names(self) -> List[str]:
         r"""
@@ -3113,12 +3120,7 @@ def display_names(self) -> List[str]:
         >>> AxisCollection([a, b, c, d]).display_names
         ['a', 'b*', '{2}', '{3}*']
         """
-        def display_name(i, axis):
-            # str(axis.name) because name can be an integer
-            name = str(axis.name) if axis.name is not None else f'{{{i}}}'
-            return (name + '*') if axis.iswildcard else name
-
-        return [display_name(i, axis) for i, axis in enumerate(self._list)]
+        return [self._display_name(axis, i) for i, axis in enumerate(self._list)]
 
     @property
     def ids(self) -> List[Union[str, int]]:
@@ -3233,10 +3235,8 @@ def info(self) -> str:
          sex [2]: 'M' 'F'
          time [4]: 2007 2008 2009 2010
         """
-        lines = [f" {name} [{len(axis)}]: {axis.labels_summary()}"
-                 for name, axis in zip(self.display_names, self._list)]
-        shape = " x ".join(str(s) for s in self.shape)
-        return ReprString('\n'.join([shape] + lines))
+        shape_str = " x ".join(str(s) for s in self.shape)
+        return ReprString(f"{shape_str}\n{self._axes_summary()}")
 
     # XXX: instead of front_if_spread, we might want to require axes to be contiguous
     #      (ie the caller would have to transpose axes before calling this)
diff --git a/larray/tests/test_array.py b/larray/tests/test_array.py
@@ -620,29 +620,49 @@ def test_getitem_guess_axis(array):
     assert_array_equal(array[g], raw[..., [0, 4, 8]])
 
     # key with duplicate axes
-    with pytest.raises(ValueError, match="key has several values for axis: age"):
+    with must_raise(ValueError, """key has several values for axis: age
+key: ([1, 2], [3, 4])"""):
         _ = array[[1, 2], [3, 4]]
 
     # key with invalid label (ie label not found on any axis)
-    with pytest.raises(ValueError, match="999 is not a valid label for any axis"):
+    with must_raise(ValueError, """999 is not a valid label for any axis:
+ age [116]: 0 1 2 ... 113 114 115
+ geo [44]: 'A11' 'A12' 'A13' ... 'A92' 'A93' 'A21'
+ sex [2]: 'M' 'F'
+ lipro [15]: 'P01' 'P02' 'P03' ... 'P13' 'P14' 'P15'"""):
         _ = array[[1, 2], 999]
 
     # key with invalid label list (ie list of labels not found on any axis)
-    with pytest.raises(ValueError, match=r"\[998, 999\] is not a valid label for any axis"):
+    with must_raise(ValueError, """[998, 999] is not a valid label for any axis:
+ age [116]: 0 1 2 ... 113 114 115
+ geo [44]: 'A11' 'A12' 'A13' ... 'A92' 'A93' 'A21'
+ sex [2]: 'M' 'F'
+ lipro [15]: 'P01' 'P02' 'P03' ... 'P13' 'P14' 'P15'"""):
         _ = array[[1, 2], [998, 999]]
 
     # key with partial invalid list (ie list containing a label not found
     # on any axis)
-    # FIXME: the message should be the same as for 999, 4 (ie it should NOT mention age).
-    with pytest.raises(ValueError, match=r"age\[3, 999\] is not a valid label for any axis"):
+    with must_raise(ValueError, "age[3, 999] is not a valid label for the 'age' axis with labels: 0, 1, 2, 3, 4, 5, 6, "
+                                "7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, "
+                                "29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, "
+                                "50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, "
+                                "71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, "
+                                "92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, "
+                                "110, 111, 112, 113, 114, 115"):
         _ = array[[1, 2], [3, 999]]
 
-    with pytest.raises(ValueError, match=r"\[999, 4\] is not a valid label for any axis"):
+    with must_raise(ValueError, """[999, 4] is not a valid label for any axis:
+ age [116]: 0 1 2 ... 113 114 115
+ geo [44]: 'A11' 'A12' 'A13' ... 'A92' 'A93' 'A21'
+ sex [2]: 'M' 'F'
+ lipro [15]: 'P01' 'P02' 'P03' ... 'P13' 'P14' 'P15'"""):
         _ = array[[1, 2], [999, 4]]
 
     # ambiguous key
     arr = ndtest("a=l0,l1;b=l1,l2")
-    with pytest.raises(ValueError, match=r"l1 is ambiguous \(valid in a, b\)"):
+    with must_raise(ValueError, """'l1' is ambiguous, it is valid in the following axes:
+ a [2]: 'l0' 'l1'
+ b [2]: 'l1' 'l2'"""):
         _ = arr['l1']
 
     # ambiguous key disambiguated via string
@@ -1219,7 +1239,7 @@ def test_points_indexer_setitem():
 
     arr = ndtest(2)
     # XXX: we might want to raise KeyError or IndexError instead?
-    with must_raise(ValueError, match="'b1' is not a valid label for any axis"):
+    with must_raise(ValueError, "'b1' is not a valid label for any axis:\n a [2]: 'a0' 'a1'"):
         arr.points['a0', 'b1'] = 42
 
     # test when broadcasting is involved
@@ -1943,6 +1963,14 @@ def test_group_agg_guess_axis(array):
     arr = ndtest(4)
     assert arr.sum('a3,a1') == 4
 
+    # ambiguous label and anonymous axes
+    arr = ndtest([Axis("b1,b2"), Axis("b0..b2")])
+    msg = """'b1' is ambiguous, it is valid in the following axes:
+ {0} [2]: 'b1' 'b2'
+ {1} [3]: 'b0' 'b1' 'b2'"""
+    with must_raise(ValueError, msg=msg):
+        arr.sum('b1;b0,b1')
+
 
 def test_group_agg_label_group(array):
     age, geo, sex, lipro = array.axes
@@ -2423,7 +2451,8 @@ def test_sum_with_groups_from_other_axis(small_array):
     # use a group (from another axis) which is incompatible with the axis of
     # the same name in the array
     lipro4 = Axis('lipro=P01,P03,P16')
-    with pytest.raises(ValueError, match=r"lipro\['P01', 'P16'\] is not a valid label for any axis"):
+    codes = "'P01', 'P02', 'P03', 'P04', 'P05', 'P06', 'P07', 'P08', 'P09', 'P10', 'P11', 'P12', 'P13', 'P14', 'P15'"
+    with must_raise(ValueError, f"lipro['P01', 'P16'] is not a valid label for the 'lipro' axis with labels: {codes}"):
         small_array.sum(lipro4['P01,P16'])
 
 
diff --git a/larray/tests/test_axiscollection.py b/larray/tests/test_axiscollection.py
@@ -370,5 +370,14 @@ def test_repr(col):
 ])"""
 
 
+def test_setlabels():
+    # test when the label is ambiguous AND the axes are anonymous
+    axes = AxisCollection([Axis("b1,b2"), Axis("b0..b2")])
+    with must_raise(ValueError, msg="""'b1' is ambiguous, it is valid in the following axes:
+ {0} [2]: 'b1' 'b2'
+ {1} [3]: 'b0' 'b1' 'b2'"""):
+        axes.set_labels({'b1': 'b_one'})
+
+
 if __name__ == "__main__":
     pytest.main()
