FIX: fix memory leak in swt_axis Cython routine.

previously some output_info.shape, output_info.strides pointers could
potentially point to a garbage-collected object.

Also changed cA, cD to initialize as empty rather than zeros for efficiency.

Author: grlee77

pywt/_extensions/_swt.pyx

@@ -102,8 +102,7 @@ cpdef swt_axis(np.ndarray data, Wavelet wavelet, size_t level,
     # memory-views do not support n-dimensional arrays, use np.ndarray instead
     cdef common.ArrayInfo data_info, output_info
     cdef np.ndarray cD, cA
-    # Explicit input_shape necessary to prevent memory leak
-    cdef size_t[::1] input_shape, output_shape
+    cdef size_t[::1] output_shape
     cdef size_t end_level = start_level + level
     cdef int i, retval
 
@@ -122,28 +121,23 @@ cpdef swt_axis(np.ndarray data, Wavelet wavelet, size_t level,
         raise ValueError(msg)
 
     data = data.astype(_check_dtype(data), copy=False)
-
-    input_shape = <size_t [:data.ndim]> <size_t *> data.shape
-    output_shape = input_shape.copy()
-    output_shape[axis] = common.swt_buffer_length(data.shape[axis])
-    if output_shape[axis] != input_shape[axis]:
-        raise RuntimeError("swt_axis assumes output_shape is the same as "
-                           "input_shape")
+    # For SWT, the output matches the shape of the input
+    output_shape = <size_t [:data.ndim]> <size_t *> data.shape
 
     data_info.ndim = data.ndim
     data_info.strides = <pywt_index_t *> data.strides
     data_info.shape = <size_t *> data.shape
 
-    cA = np.empty(output_shape, data.dtype)
-    output_info.ndim = cA.ndim
-    output_info.strides = <pywt_index_t *> cA.strides
-    output_info.shape = <size_t *> cA.shape
+    output_info.ndim = data.ndim
 
     ret = []
     for i in range(start_level+1, end_level+1):
-
+        cA = np.empty(output_shape, dtype=data.dtype)
+        cD = np.empty(output_shape, dtype=data.dtype)
+        # strides won't match data_info.strides if data is not C-contiguous
+        output_info.strides = <pywt_index_t *> cA.strides
+        output_info.shape = <size_t *> cA.shape
         if data.dtype == np.float64:
-            cA = np.zeros(output_shape, dtype=np.float64)
             with nogil:
                 retval = c_wt.double_downcoef_axis(
                     <double *> data.data, data_info,
@@ -152,8 +146,8 @@ cpdef swt_axis(np.ndarray data, Wavelet wavelet, size_t level,
                     common.COEF_APPROX, common.MODE_PERIODIZATION,
                     i, common.SWT_TRANSFORM)
             if retval:
-                raise RuntimeError("C wavelet transform failed")
-            cD = np.zeros(output_shape, dtype=np.float64)
+                raise RuntimeError(
+                    "C wavelet transform failed with error code %d" % retval)
             with nogil:
                 retval = c_wt.double_downcoef_axis(
                     <double *> data.data, data_info,
@@ -162,9 +156,9 @@ cpdef swt_axis(np.ndarray data, Wavelet wavelet, size_t level,
                     common.COEF_DETAIL, common.MODE_PERIODIZATION,
                     i, common.SWT_TRANSFORM)
             if retval:
-                raise RuntimeError("C wavelet transform failed")
+                raise RuntimeError(
+                    "C wavelet transform failed with error code %d" % retval)
         elif data.dtype == np.float32:
-            cA = np.zeros(output_shape, dtype=np.float32)
             with nogil:
                 retval = c_wt.float_downcoef_axis(
                     <float *> data.data, data_info,
@@ -173,8 +167,8 @@ cpdef swt_axis(np.ndarray data, Wavelet wavelet, size_t level,
                     common.COEF_APPROX, common.MODE_PERIODIZATION,
                     i, common.SWT_TRANSFORM)
             if retval:
-                raise RuntimeError("C wavelet transform failed")
-            cD = np.zeros(output_shape, dtype=np.float32)
+                raise RuntimeError(
+                    "C wavelet transform failed with error code %d" % retval)
             with nogil:
                 retval = c_wt.float_downcoef_axis(
                     <float *> data.data, data_info,
@@ -183,15 +177,18 @@ cpdef swt_axis(np.ndarray data, Wavelet wavelet, size_t level,
                     common.COEF_DETAIL, common.MODE_PERIODIZATION,
                     i, common.SWT_TRANSFORM)
             if retval:
-                raise RuntimeError("C wavelet transform failed")
+                raise RuntimeError(
+                    "C wavelet transform failed with error code %d" % retval)
         else:
             raise TypeError("Array must be floating point, not {}"
                             .format(data.dtype))
         ret.append((cA, cD))
 
         # previous approx coeffs are the data for the next level
         data = cA
-        data_info = output_info
+        # update data_info to match the new data array
+        data_info.strides = <pywt_index_t *> data.strides
+        data_info.shape = <size_t *> data.shape
 
     ret.reverse()
     return ret