Merge pull request #254 from isuruf/slice

Fix slicing in matrices

Author: isuruf

symengine/lib/symengine_wrapper.pyx

@@ -3164,7 +3164,7 @@ cdef class DenseMatrixBase(MatrixBase):
                     row_iter = [item[0]]
 
                 if isinstance(item[1], slice):
-                    col_iter = range(*item[1].indices(self.rows))
+                    col_iter = range(*item[1].indices(self.cols))
                 elif is_sequence(item[1]):
                     col_iter = item[1]
                 else:
@@ -3213,7 +3213,7 @@ cdef class DenseMatrixBase(MatrixBase):
                 row_iter = [key[0]]
 
             if isinstance(key[1], slice):
-                col_iter = range(*key[1].indices(self.rows))
+                col_iter = range(*key[1].indices(self.cols))
             elif is_sequence(key[1]):
                 col_iter = key[1]
             else:

symengine/tests/test_matrices.py

@@ -73,6 +73,14 @@ def test_get_item():
     raises(IndexError, lambda: A[1:3, 3])
     raises(IndexError, lambda: A[1:3, -4])
 
+    A = zeros(3, 4)
+    assert list(A[0, :]) == [0, 0, 0, 0]
+    assert list(A[:, 0]) == [0, 0, 0]
+
+    A = zeros(4, 3)
+    assert list(A[:, 0]) == [0, 0, 0, 0]
+    assert list(A[0, :]) == [0, 0, 0]
+
 
 def test_set_item():
     A = DenseMatrix(3, 3, [1, 2, 3, 4, 5, 6, 7, 8, 9])
@@ -129,6 +137,16 @@ def test_set_item():
     B[[0, 1], [0, 2]] = -1
     assert B == DenseMatrix(3, 3, [-1, 2, -1, -1, 14, -1, 4, 5, 6])
 
+    A = zeros(3, 4)
+    B = ones(1, 4)
+    A[0, :] = B
+    assert A[0, :] == B
+
+    A = zeros(3, 4)
+    B = ones(3, 1)
+    A[:, 0] = B
+    assert A[:, 0] == B
+
 
 def test_set():
     i7 = Integer(7)