added support for ellipsis in __(set|get)item__

gdementen · gdementen · commit d83fbf6ddd31 · 2015-01-28T12:56:54.000+01:00
diff --git a/larray/core.py b/larray/core.py
@@ -880,6 +880,16 @@ def full_key(self, key):
         if not isinstance(key, (tuple, dict)):
             key = (key,)
 
+        # handle keys containing an Ellipsis
+        if isinstance(key, tuple):
+            num_ellipses = key.count(Ellipsis)
+            if num_ellipses > 1:
+                raise ValueError("cannot use more than one Ellipsis (...)")
+            elif num_ellipses == 1:
+                pos = key.index(Ellipsis)
+                none_slices = (slice(None),) * (self.ndim - len(key) + 1)
+                key = key[:pos] + none_slices + key[pos + 1:]
+
         # handle keys containing ValueGroups (at potentially wrong places)
         if any(isinstance(axis_key, ValueGroup) for axis_key in key):
             #XXX: support ValueGroup without axis?
diff --git a/larray/tests/test_la.py b/larray/tests/test_la.py
@@ -596,6 +596,12 @@ def test_getitem(self):
         # single None slice
         self._assert_equal_raw(la[:], raw)
 
+        # only Ellipsis
+        self._assert_equal_raw(la[...], raw)
+
+        # Ellipsis and VG
+        self._assert_equal_raw(la[..., lipro159], raw[..., [0, 4, 8]])
+
         # key with duplicate axes
         # la[[1, 5, 9], age['1,5,9']]
         self.assertRaises(ValueError, la.__getitem__, ([1, 5], age['1,5']))
@@ -650,6 +656,19 @@ def test_setitem_larray(self):
         raw[[1, 5, 9]] = raw[[2, 7, 3]] + 27.0
         self._assert_equal_raw(la, raw)
 
+        # 3) using ellipsis keys
+        # only Ellipsis
+        la = self.larray.copy()
+        la[...] = 0
+        self._assert_equal_raw(la, np.zeros_like(raw))
+
+        # Ellipsis and VG
+        la = self.larray.copy()
+        raw = self.array.copy()
+        la[..., lipro['P01,P05,P09']] = 0
+        raw[..., [0, 4, 8]] = 0
+        self._assert_equal_raw(la, raw)
+
     def test_setitem_ndarray(self):
         """
         tests LArray.__setitem__(key, value) where value is a raw ndarray.
