ENH: lib: break reference cycle in NpzFile (numpy#2048)

pv · charris · commit c4482f56b347 · 2012-02-16T19:39:29.000-07:00
This allows these objects to be freed by refcount, rather than requiring
the gc, which can be useful in some situations.
diff --git a/numpy/lib/npyio.py b/numpy/lib/npyio.py
@@ -10,6 +10,7 @@
 import sys
 import itertools
 import warnings
+import weakref
 from operator import itemgetter
 
 from cPickle import load as _cload, loads
@@ -108,7 +109,8 @@ class BagObj(object):
 
     """
     def __init__(self, obj):
-        self._obj = obj
+        # Use weakref to make NpzFile objects collectable by refcount
+        self._obj = weakref.proxy(obj)
     def __getattribute__(self, key):
         try:
             return object.__getattribute__(self, '_obj')[key]
@@ -212,6 +214,7 @@ def close(self):
         if self.fid is not None:
             self.fid.close()
             self.fid = None
+        self.f = None # break reference cycle
 
     def __del__(self):
         self.close()
diff --git a/numpy/lib/tests/test_io.py b/numpy/lib/tests/test_io.py
@@ -6,6 +6,7 @@
 import time
 from datetime import datetime
 import warnings
+import gc
 from numpy.testing.utils import WarningManager
 
 import numpy as np
@@ -1525,6 +1526,20 @@ def test_npzfile_dict():
 
     assert_('x' in list(z.iterkeys()))
 
+def test_load_refcount():
+    # Check that objects returned by np.load are directly freed based on
+    # their refcount, rather than needing the gc to collect them.
+
+    f = StringIO()
+    np.savez(f, [1, 2, 3])
+    f.seek(0)
+
+    gc.collect()
+    n_before = len(gc.get_objects())
+    np.load(f)
+    n_after = len(gc.get_objects())
+
+    assert_equal(n_before, n_after)
 
 if __name__ == "__main__":
     run_module_suite()
