Merge pull request #10622 from ahaldane/fix_arrayprint_recursive_closure_backport

BUG: deallocate recursive closure in arrayprint.py (1.14 backport)

2 files changed, 103 insertions, 88 deletions

diff --git a/numpy/core/arrayprint.py b/numpy/core/arrayprint.py
index 284c57867..1f193ecfc 100644
--- a/numpy/core/arrayprint.py
+++ b/numpy/core/arrayprint.py
@@ -642,118 +642,121 @@ def _extendLine(s, line, word, line_width, next_line_prefix, legacy):
     line += word
     return s, line
 
+def _recursive_fmt(param, index, indent, curr_width):
+    """
+    Helper function for _formatArray, to recursively print array elements.
 
-def _formatArray(a, format_function, line_width, next_line_prefix,
-                 separator, edge_items, summary_insert, legacy):
-    """formatArray is designed for two modes of operation:
+    "param" are the values that are invariant under recursion, including
+    the array to be printed itself. index, indent and curr_width
+    are updated during recursion.
+    """
+    # unpack parameters
+    a, format_function, separator, edge_items, summary_insert, legacy = param
 
-    1. Full output
+    axis = len(index)
+    axes_left = a.ndim - axis
 
-    2. Summarized output
+    if axes_left == 0:
+        return format_function(a[index])
 
-    """
-    def recurser(index, hanging_indent, curr_width):
-        """
-        By using this local function, we don't need to recurse with all the
-        arguments. Since this function is not created recursively, the cost is
-        not significant
-        """
-        axis = len(index)
-        axes_left = a.ndim - axis
+    # when recursing, add a space to align with the [ added, and reduce the
+    # length of the line by 1
+    next_indent = indent + ' '
+    if legacy == '1.13':
+        next_width = curr_width
+    else:
+        next_width = curr_width - len(']')
 
-        if axes_left == 0:
-            return format_function(a[index])
+    a_len = a.shape[axis]
+    show_summary = summary_insert and 2*edge_items < a_len
+    if show_summary:
+        leading_items = edge_items
+        trailing_items = edge_items
+    else:
+        leading_items = 0
+        trailing_items = a_len
 
-        # when recursing, add a space to align with the [ added, and reduce the
-        # length of the line by 1
-        next_hanging_indent = hanging_indent + ' '
-        if legacy == '1.13':
-            next_width = curr_width
-        else:
-            next_width = curr_width - len(']')
+    # stringify the array with the hanging indent on the first line too
+    s = ''
 
-        a_len = a.shape[axis]
-        show_summary = summary_insert and 2*edge_items < a_len
-        if show_summary:
-            leading_items = edge_items
-            trailing_items = edge_items
+    # last axis (rows) - wrap elements if they would not fit on one line
+    if axes_left == 1:
+        # the length up until the beginning of the separator / bracket
+        if legacy == '1.13':
+            elem_width = curr_width - len(separator.rstrip())
         else:
-            leading_items = 0
-            trailing_items = a_len
+            elem_width = curr_width - max(len(separator.rstrip()), len(']'))
 
-        # stringify the array with the hanging indent on the first line too
-        s = ''
+        line = indent
+        for i in range(leading_items):
+            word = _recursive_fmt(param, index + (i,), next_indent, next_width)
+            s, line = _extendLine(s, line, word, elem_width, indent, legacy)
+            line += separator
 
-        # last axis (rows) - wrap elements if they would not fit on one line
-        if axes_left == 1:
-            # the length up until the beginning of the separator / bracket
+        if show_summary:
+            s, line = _extendLine(
+                s, line, summary_insert, elem_width, indent, legacy)
             if legacy == '1.13':
-                elem_width = curr_width - len(separator.rstrip())
+                line += ", "
             else:
-                elem_width = curr_width - max(len(separator.rstrip()), len(']'))
-
-            line = hanging_indent
-            for i in range(leading_items):
-                word = recurser(index + (i,), next_hanging_indent, next_width)
-                s, line = _extendLine(
-                    s, line, word, elem_width, hanging_indent, legacy)
                 line += separator
 
-            if show_summary:
-                s, line = _extendLine(
-                    s, line, summary_insert, elem_width, hanging_indent, legacy)
-                if legacy == '1.13':
-                    line += ", "
-                else:
-                    line += separator
-
-            for i in range(trailing_items, 1, -1):
-                word = recurser(index + (-i,), next_hanging_indent, next_width)
-                s, line = _extendLine(
-                    s, line, word, elem_width, hanging_indent, legacy)
-                line += separator
+        for i in range(trailing_items, 1, -1):
+            word = _recursive_fmt(param, index + (-i,), next_indent, next_width)
+            s, line = _extendLine(s, line, word, elem_width, indent, legacy)
+            line += separator
+
+        if legacy == '1.13':
+            # width of the separator is not considered on 1.13
+            elem_width = curr_width
+        word =_recursive_fmt(param, index + (-1,), next_indent, next_width)
+        s, line = _extendLine(
+            s, line, word, elem_width, indent, legacy)
+
+        s += line
+
+    # other axes - insert newlines between rows
+    else:
+        s = ''
+        line_sep = separator.rstrip() + '\n'*(axes_left - 1)
+
+        for i in range(leading_items):
+            nested =_recursive_fmt(param, index + (i,), next_indent, next_width)
+            s += indent + nested + line_sep
 
+        if show_summary:
             if legacy == '1.13':
-                # width of the seperator is not considered on 1.13
-                elem_width = curr_width
-            word = recurser(index + (-1,), next_hanging_indent, next_width)
-            s, line = _extendLine(
-                s, line, word, elem_width, hanging_indent, legacy)
+                # trailing space, fixed nbr of newlines, and fixed separator
+                s += indent + summary_insert + ", \n"
+            else:
+                s += indent + summary_insert + line_sep
 
-            s += line
+        for i in range(trailing_items, 1, -1):
+            nested = _recursive_fmt(param, index + (-i,), next_indent,
+                                    next_width)
+            s += indent + nested + line_sep
 
-        # other axes - insert newlines between rows
-        else:
-            s = ''
-            line_sep = separator.rstrip() + '\n'*(axes_left - 1)
+        nested =_recursive_fmt(param, index + (-1,), next_indent, next_width)
+        s += indent + nested
 
-            for i in range(leading_items):
-                nested = recurser(index + (i,), next_hanging_indent, next_width)
-                s += hanging_indent + nested + line_sep
+    # remove the hanging indent, and wrap in []
+    s = '[' + s[len(indent):] + ']'
+    return s
 
-            if show_summary:
-                if legacy == '1.13':
-                    # trailing space, fixed nbr of newlines, and fixed separator
-                    s += hanging_indent + summary_insert + ", \n"
-                else:
-                    s += hanging_indent + summary_insert + line_sep
+def _formatArray(a, format_function, line_width, next_line_prefix,
+                 separator, edge_items, summary_insert, legacy):
+    """_formatArray is designed for two modes of operation:
 
-            for i in range(trailing_items, 1, -1):
-                nested = recurser(index + (-i,), next_hanging_indent, next_width)
-                s += hanging_indent + nested + line_sep
+    1. Full output
 
-            nested = recurser(index + (-1,), next_hanging_indent, next_width)
-            s += hanging_indent + nested
+    2. Summarized output
 
-        # remove the hanging indent, and wrap in []
-        s = '[' + s[len(hanging_indent):] + ']'
-        return s
+    """
 
     # invoke the recursive part with an initial index and prefix
-    return recurser(
-        index=(),
-        hanging_indent=next_line_prefix,
-        curr_width=line_width)
+    param = a, format_function, separator, edge_items, summary_insert, legacy
+    return _recursive_fmt(param, index=(), indent=next_line_prefix,
+                          curr_width=line_width)
 
 
 def _none_or_positive_arg(x, name):
diff --git a/numpy/core/tests/test_arrayprint.py b/numpy/core/tests/test_arrayprint.py
index f70b6a333..7e76d442d 100644
--- a/numpy/core/tests/test_arrayprint.py
+++ b/numpy/core/tests/test_arrayprint.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 from __future__ import division, absolute_import, print_function
 
-import sys
+import sys, gc
 
 import numpy as np
 from numpy.testing import (
@@ -355,6 +355,18 @@ class TestArray2String(object):
             "[ 'xxxxx']"
         )
 
+    def test_refcount(self):
+        # make sure we do not hold references to the array due to a recursive
+        # closure (gh-10620)
+        gc.disable()
+        a = np.arange(2)
+        r1 = sys.getrefcount(a)
+        np.array2string(a)
+        np.array2string(a)
+        r2 = sys.getrefcount(a)
+        gc.collect()
+        gc.enable()
+        assert_(r1 == r2)
 
 class TestPrintOptions(object):
     """Test getting and setting global print options."""