BUG: ensure integer shape, strides for as_strided

mbyt · mbyt · commit d34a3f096b61 · 2014-11-25T21:52:20.000+01:00
diff --git a/lib/matplotlib/mlab.py b/lib/matplotlib/mlab.py
@@ -590,6 +590,19 @@ def stride_windows(x, n, noverlap=None, axis=0):
     if n > x.size:
         raise ValueError('n cannot be greater than the length of x')
 
+    # as_strided can lead to memory out of bound access for non
+    # integer noverlap or n
+    if not isinstance(noverlap, (int, long)):
+        warnings.warn("Using a non integer 'noverlap' instead of an integer "
+                      "will result in an error in the future",
+                      cbook.mplDeprecation)
+        noverlap = int(round(noverlap))
+    if not isinstance(n, (int, long)):
+        warnings.warn("Using a non integer 'n' instead of an integer "
+                      "will result in an error in the future",
+                      cbook.mplDeprecation)
+        n = int(round(n))
+
     step = n - noverlap
     if axis == 0:
         shape = (n, (x.shape[-1]-noverlap)//step)
@@ -642,6 +655,14 @@ def stride_repeat(x, n, axis=0):
     if n < 1:
         raise ValueError('n cannot be less than 1')
 
+    # as_strided might lead to memory out of bound access for non
+    # integer n
+    if not isinstance(n, (int, long)):
+        warnings.warn("Using a non integer 'n' instead of an integer "
+                      "will result in an error in the future",
+                      cbook.mplDeprecation)
+        n = int(round(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
@@ -4,6 +4,7 @@
 import six
 
 import tempfile
+import warnings
 
 from numpy.testing import assert_allclose, assert_array_equal
 import numpy.ma.testutils as matest
@@ -283,6 +284,22 @@ 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[:] = np.sin(np.arange(N)/N * 2 * np.pi * n_sin_periods)
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', cbook.mplDeprecation)
+            y_strided = mlab.stride_windows(y, n=330, noverlap=0.6)
+            assert_equal(sum(np.isnan(y_strided.flat[:])), 0)
+            y_strided = mlab.stride_windows(y, n=333.3, noverlap=0)
+            assert_equal(sum(np.isnan(y_strided.flat[:])), 0)
+            y_strided = mlab.stride_repeat(y, n=33.815)
+            assert_equal(sum(np.isnan(y_strided.flat[:])), 0)
+
 
 class csv_testcase(CleanupTestCase):
     def setUp(self):
