BUG: ensure integer shape, strides for as_strided

mbyt · mbyt · commit 43fd6bea336e · 2014-11-26T20:17:44.000+01:00
diff --git a/lib/matplotlib/mlab.py b/lib/matplotlib/mlab.py
@@ -590,6 +590,11 @@ def stride_windows(x, n, noverlap=None, axis=0):
     if n > x.size:
         raise ValueError('n cannot be greater than the length of x')
 
+    # np.lib.stride_tricks.as_strided easily leads to memory corruption for
+    # non integer shape and strides, i.e. noverlap or n. See #3845.
+    noverlap = int(noverlap)
+    n = int(n)
+
     step = n - noverlap
     if axis == 0:
         shape = (n, (x.shape[-1]-noverlap)//step)
@@ -642,6 +647,10 @@ def stride_repeat(x, n, axis=0):
     if n < 1:
         raise ValueError('n cannot be less than 1')
 
+    # np.lib.stride_tricks.as_strided easily leads to memory corruption for
+    # non integer shape and strides, i.e. n. See #3845.
+    n = int(n)
+
     if axis == 0:
         shape = (n, x.size)
         strides = (0, x.strides[0])
diff --git a/lib/matplotlib/tests/test_mlab.py b/lib/matplotlib/tests/test_mlab.py
@@ -283,6 +283,25 @@ def test_stride_windows_n32_noverlap0_axis1_unflatten(self):
         assert_equal(y.shape, x1.shape)
         assert_array_equal(y, x1)
 
+    def test_stride_ensure_integer_type(self):
+        N = 1000
+        n_sin_periods = 10
+        x = np.empty(3*N, dtype='>f4')
+        x.fill(np.NaN)
+        y = x[N:2*N]
+        y.fill(0.3)
+        # previous to #3845 lead to corrupt access
+        y_strided = mlab.stride_windows(y, n=330, noverlap=0.6)
+        assert_array_equal(y_strided, 0.3)
+        # previous to #3845 lead to corrupt access
+        y_strided = mlab.stride_windows(y, n=333.3, noverlap=0)
+        assert_array_equal(y_strided, 0.3)
+        # even previous to #3845 could not find any problematic
+        # configuration however, let's be sure it's not accidentally
+        # introduced
+        y_strided = mlab.stride_repeat(y, n=33.815)
+        assert_array_equal(y_strided, 0.3)
+
 
 class csv_testcase(CleanupTestCase):
     def setUp(self):
