Merge pull request #429 from 87/fix_insert

njsmith · njsmith · commit aed8fc547bfa · 2012-09-11T11:07:32.000-07:00
Fix for issues #392 and #378
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
@@ -3591,19 +3591,21 @@ def insert(arr, obj, values, axis=None):
     slobj = [slice(None)]*ndim
     N = arr.shape[axis]
     newshape = list(arr.shape)
-    if isinstance(obj, (int, long, integer)):
 
+    if isinstance(obj, (int, long, integer)):
         if (obj < 0): obj += N
         if obj < 0 or obj > N:
             raise ValueError(
                     "index (%d) out of range (0<=index<=%d) "\
                     "in dimension %d" % (obj, N, axis))
-          
-        if isinstance(values, (int, long, integer)):            
-            obj = [obj]           
+        if isscalar(values):
+            obj = [obj]
         else:
-            obj = [obj] * len(values) 
-
+            values = asarray(values)
+            if ndim > values.ndim:
+                obj = [obj]
+            else:
+                obj = [obj] * len(values)
 
     elif isinstance(obj, slice):
         # turn it into a range object
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
@@ -147,6 +147,15 @@ def test_basic(self):
         assert_equal(insert(a, [1, 1, 1], [1, 2, 3]), [1, 1, 2, 3, 2, 3])
         assert_equal(insert(a, 1,[1,2,3]), [1, 1, 2, 3, 2, 3])
         assert_equal(insert(a,[1,2,3],9),[1,9,2,9,3,9])
+        b = np.array([0, 1], dtype=np.float64)
+        assert_equal(insert(b, 0, b[0]), [0., 0., 1.])
+    def test_multidim(self):
+        a = [[1, 1, 1]]
+        r = [[2, 2, 2],
+             [1, 1, 1]]
+        assert_equal(insert(a, 0, [2, 2, 2], axis=0), r)
+        assert_equal(insert(a, 0, 2, axis=0), r)
+        assert_equal(insert(a, 2, 2, axis=1), [[1, 1, 2, 1]])
 
 class TestAmax(TestCase):
     def test_basic(self):
