FIX: raise error on varying number of separators in Array.split_axes (closes #1089)

gdementen · gdementen · commit c2080f2a51d6 · 2023-10-23T16:02:41.000+02:00
diff --git a/doc/source/changes/version_0_34_2.rst.inc b/doc/source/changes/version_0_34_2.rst.inc
@@ -58,3 +58,7 @@ Fixes
 * fixed Array.reindex when using an axis object from the array as `axes_to_reindex` (closes :issue:`1088`).
 
 * fixed Array.reindex({axis: list_of_labels}) (closes :issue:`1068`).
+
+* Array.split_axes now raises an explicit error when some labels contain
+  more separators than others, instead of silently dropping part of those
+  labels, or even some data (closes :issue:`1089`).
diff --git a/larray/core/axis.py b/larray/core/axis.py
@@ -382,19 +382,30 @@ def split(self, sep='_', names=None, regex=None, return_labels=False) \
         if not regex:
             # np.char.split does not work on arrays with object dtype
             labels = self.labels if self.labels.dtype.kind != 'O' else self.labels.astype(str)
-            # gives us an array of lists
+            # split_labels is an array of lists
             split_labels = np.char.split(labels, sep)
         else:
             match = re.compile(regex).match
+            # split_labels is a list of tuples
             split_labels = [match(label).groups() for label in self.labels]
+        first_split_label_length = len(split_labels[0])
+        # TODO: when our lowest supported version will be Python 3.10, we should use
+        #       strict=True instead of checking lengths explicitly
+        if any(len(split_label) != first_split_label_length
+               for split_label in split_labels):
+            raise ValueError("not all labels have the same number of separators")
+        indexing_labels = tuple(zip(*split_labels))
         if names is None:
-            names = [None] * len(split_labels)
-        indexing_labels = zip(*split_labels)
-        if return_labels:
-            indexing_labels = tuple(indexing_labels)
+            names = [None] * first_split_label_length
+        num_axes = len(indexing_labels)
+        if num_axes != len(names):
+            raise ValueError(f"number of resulting axes ({num_axes}) differs "
+                             f"from number of resulting axes names "
+                             f"({len(names)})")
         # not using np.unique because we want to keep the original order
         split_axes = [Axis(unique_list(ax_labels), name) for ax_labels, name in zip(indexing_labels, names)]
         if return_labels:
+            assert len(split_axes) == num_axes
             indexing_labels = tuple(axis[labels] for axis, labels in zip(split_axes, indexing_labels))
             return split_axes, indexing_labels
         else:
diff --git a/larray/tests/test_array.py b/larray/tests/test_array.py
@@ -5497,6 +5497,20 @@ def test_split_axes():
     # when no axis is specified and no axis contains the sep, split_axes is a no-op.
     assert_larray_equal(combined.split_axes(), combined)
 
+    # with varying sep characters in labels (issue #1089)
+    arr = ndtest("a_b=a0_b0,a0_b1_1,a0_b1_2")
+    # a_b  a0_b0  a0_b1_1  a0_b1_2
+    #          0        1        2
+    with must_raise(ValueError, "not all labels have the same number of separators"):
+        arr.split_axes()
+
+    # with different number of sep characters in labels than in axis name
+    arr = ndtest("a_b=a0_b0_1,a0_b1_1,a0_b1_2")
+    # a_b  a0_b0_1  a0_b1_1  a0_b1_2
+    #            0        1        2
+    with must_raise(ValueError, "number of resulting axes (3) differs from number of resulting axes names (2)"):
+        arr.split_axes()
+
 
 def test_stack():
     # stack along a single axis
