BUG fix 1d bool larry indexing

kwgoodman · Jun 29, 2012 · 706670a · 706670a
1 parent bbf7c96
commit 706670a
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 9 deletions.
diff --git a/README.rst b/README.rst
@@ -122,9 +122,9 @@ After you have installed ``la``, run the suite of unit tests::
     >>> import la
     >>> la.test()
     <snip>
-    Ran 3014 tests in 9.225s
+    Ran 3018 tests in 9.225s
     OK
-    <nose.result.TextTestResult run=3014 errors=0 failures=0>
+    <nose.result.TextTestResult run=3018 errors=0 failures=0>
     
 The ``la`` package contains C extensions that speed up common alignment
 operations such as adding two unaligned larrys. If the C extensions don't

diff --git a/la/deflarry.py b/la/deflarry.py
@@ -1685,11 +1685,13 @@ def __getitem__(self, index):
                         allscalar = False
                     elif typ is larry:
                         if idx.dtype.type != np.bool_:
-                            raise ValueError("If index is larry it must of bool dtype")
+                            raise IndexError("If index is larry it must of bool dtype")
                         if idx.ndim != 1:
-                            raise ValueError("If index is larry it must be 1d")
+                            raise IndexError("If index is larry it must be 1d")
+                        if self.label[ax] != idx.label[0]:
+                            raise IndexError("When indexing larrys must be aligned")
                         try:
-                            lab = [self.label[0][j] for j, z in
+                            lab = [self.label[ax][j] for j, z in
                                                         enumerate(idx.x) if z]
                         except IndexError:
                             raise IndexError, 'index out of range'
@@ -1746,9 +1748,11 @@ def __getitem__(self, index):
             label[0] = lab                 
         elif typidx is larry:
             if index.dtype.type != np.bool_:
-                raise ValueError("If index is larry it must of bool dtype")
+                raise IndexError("If index is larry it must of bool dtype")
             if index.ndim != 1:
-                raise ValueError("If index is larry it must be 1d")
+                raise IndexError("If index is larry it must be 1d")
+            if self.label[0] != index.label[0]:
+                raise IndexError("When indexing larrys must be aligned")
             try:
                 lab = [self.label[0][j] for j, z in
                                             enumerate(index) if z]

diff --git a/la/tests/deflarry_getitem_test.py b/la/tests/deflarry_getitem_test.py
@@ -230,7 +230,7 @@ def test_getitem_26():
     idx1 = larry([True, False, True])
     actual = lar1[idx1]
     arr = lar1.x[idx1.x]
-    assert_equal(arr.shape, actual.shape)
+    assert_array_equal(arr, actual.x)
     idx2 = [0, 2]
     desired = lar1[idx2]
     ale(actual, desired, original=lar1)
@@ -241,7 +241,58 @@ def test_getitem_27():
     idx1 = larry([True, False])
     actual = lar1[:,idx1]
     arr = lar1.x[:,idx1.x]
-    assert_equal(arr.shape, actual.shape)
+    assert_array_equal(arr, actual.x)
+    idx2 = [0]
+    desired = lar1[:,idx2]
+    ale(actual, desired, original=lar1)
+
+def test_getitem_28():
+    "larry.getitem #28"
+    lar1, lar2 = make_larrys()
+    idx1 = larry([True, False, True, False])
+    actual = lar2[idx1]
+    arr = lar2.x[idx1.x]
+    assert_array_equal(arr, actual.x)
+    idx2 = [0, 2]
+    desired = lar2[idx2]
+    ale(actual, desired, original=lar2)
+
+def test_getitem_29():
+    "larry.getitem #29"
+    lar1, lar2 = make_larrys()
+    idx = larry([True, False, True, False])
+    assert_raises(IndexError, lar1.__getitem__, idx[::-1])
+    assert_raises(IndexError, lar2.__getitem__, idx[::-1])
+    idx = larry([0, 1, 2, 3])
+    assert_raises(IndexError, lar1.__getitem__, idx)
+    assert_raises(IndexError, lar2.__getitem__, idx)
+    idx = larry([True, False])
+    assert_raises(IndexError, lar1.__getitem__, (slice(None), idx[::-1]))    
+
+def test_getitem_30():
+    "larry.getitem #30"
+    lar1, lar2 = make_larrys()
+    lar1.label[0] = ['a', 'b', 'c']
+    lar1.label[1] = ['A', 'B']
+    lar2.label[0] = ['a', 'b', 'c', 'd']
+    idx1 = larry([True, False, True], label=[['a', 'b', 'c']])
+    actual = lar1[idx1]
+    arr = lar1.x[idx1.x]
+    assert_array_equal(arr, actual.x)
+    idx2 = [0, 2]
+    desired = lar1[idx2]
+    ale(actual, desired, original=lar1)
+
+def test_getitem_31():
+    "larry.getitem #31"
+    lar1, lar2 = make_larrys()
+    lar1.label[0] = ['a', 'b', 'c']
+    lar1.label[1] = ['A', 'B']
+    lar2.label[0] = ['a', 'b', 'c', 'd']
+    idx1 = larry([True, False], label=[['A', 'B']])
+    actual = lar1[:,idx1]
+    arr = lar1.x[:,idx1.x]
+    assert_array_equal(arr, actual.x)
     idx2 = [0]
     desired = lar1[:,idx2]
     ale(actual, desired, original=lar1)