implement la[bool_array]

gdementen · gdementen · commit ce5caa37d71e · 2015-01-28T14:20:34.000+01:00
diff --git a/larray/core.py b/larray/core.py
@@ -970,6 +970,9 @@ def translated_key(self, key):
     def __getitem__(self, key, collapse_slices=False):
         data = np.asarray(self)
 
+        if isinstance(key, np.ndarray) and np.issubdtype(key.dtype, bool):
+            return data[key]
+
         translated_key = self.translated_key(self.full_key(key))
 
         axes = [axis.subaxis(axis_key)
diff --git a/larray/tests/test_la.py b/larray/tests/test_la.py
@@ -606,6 +606,15 @@ def test_getitem(self):
         # la[[1, 5, 9], age['1,5,9']]
         self.assertRaises(ValueError, la.__getitem__, ([1, 5], age['1,5']))
 
+    def test_getitem_bool_array_key(self):
+        raw = self.array
+        la = self.larray
+
+        # LArray key
+        self._assert_equal_raw(la[la < 5], raw[raw < 5])
+        # ndarray key
+        self._assert_equal_raw(la[raw < 5], raw[raw < 5])
+
     def test_setitem_larray(self):
         """
         tests LArray.__setitem__(key, value) where value is an LArray
