[pypy-commit] pypy default: refactor the autoflush of streams: we cannot keep a set of w_iobase instances, else they would be never collected by the GC. Instead, we keep a set of 'holders', which have a weakref to the actual stream. When the stream is closed, the holder is removed from the set

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: 
Changeset: r52807:4867e38be6fb
Date: 2012-02-23 16:11 +0100
http://bitbucket.org/pypy/pypy/changeset/4867e38be6fb/

Log:	refactor the autoflush of streams: we cannot keep a set of w_iobase
	instances, else they would be never collected by the GC. Instead, we
	keep a set of 'holders', which have a weakref to the actual stream.
	When the stream is closed, the holder is removed from the set

diff --git a/pypy/module/_io/__init__.py b/pypy/module/_io/__init__.py
--- a/pypy/module/_io/__init__.py
+++ b/pypy/module/_io/__init__.py
@@ -39,5 +39,6 @@
 
     def shutdown(self, space):
         # at shutdown, flush all open streams.  Ignore I/O errors.
-        from pypy.module._io.interp_iobase import flush_all_streams
-        flush_all_streams(space)
+        from pypy.module._io.interp_iobase import get_autoflushher
+        get_autoflushher(space).flush_all(space)
+
diff --git a/pypy/module/_io/interp_iobase.py b/pypy/module/_io/interp_iobase.py
--- a/pypy/module/_io/interp_iobase.py
+++ b/pypy/module/_io/interp_iobase.py
@@ -5,6 +5,8 @@
 from pypy.interpreter.gateway import interp2app
 from pypy.interpreter.error import OperationError, operationerrfmt
 from pypy.rlib.rstring import StringBuilder
+from pypy.rlib import rweakref
+
 
 DEFAULT_BUFFER_SIZE = 8192
 
@@ -43,7 +45,8 @@
         self.space = space
         self.w_dict = space.newdict()
         self.__IOBase_closed = False
-        register_flushable_stream(space, self)
+        self.streamholder = None # needed by AutoFlusher
+        get_autoflushher(space).add(self)
 
     def getdict(self, space):
         return self.w_dict
@@ -99,8 +102,7 @@
             space.call_method(self, "flush")
         finally:
             self.__IOBase_closed = True
-            flushable_streams = get_flushable_streams(space)
-            del flushable_streams[self]
+            get_autoflushher(space).remove(self)
 
     def flush_w(self, space):
         if self._CLOSED():
@@ -312,26 +314,46 @@
 # functions to make sure that all streams are flushed on exit
 # ------------------------------------------------------------
 
-class IoState:
+class StreamHolder(object):
+
+    def __init__(self, w_iobase):
+        self.w_iobase_ref = rweakref.ref(w_iobase)
+        w_iobase.autoflusher = self
+
+    def autoflush(self, space):
+        w_iobase = self.w_iobase_ref()
+        if w_iobase is not None:
+            space.call_method(w_iobase, 'flush') # XXX: ignore IOErrors?
+        
+
+class AutoFlusher(object):
+    
     def __init__(self, space):
-        self.flushable_streams = {}
+        self.streams = {}
 
-def get_flushable_streams(space):
-    return space.fromcache(IoState).flushable_streams
+    def add(self, w_iobase):
+        assert w_iobase.streamholder is None
+        holder = StreamHolder(w_iobase)
+        w_iobase.streamholder = holder
+        self.streams[holder] = None
 
-def register_flushable_stream(space, w_stream):
-    streams = get_flushable_streams(space)
-    streams[w_stream] = None
+    def remove(self, w_iobase):
+        holder = w_iobase.streamholder
+        if holder is not None:
+            del self.streams[holder]
 
-def flush_all_streams(space):
-    flushable_streams = get_flushable_streams(space)
-    while flushable_streams:
-        for w_stream in flushable_streams.keys():
-            assert isinstance(w_stream, W_IOBase)
-            try:
-                del flushable_streams[w_stream]
-            except KeyError:
-                pass    # key was removed in the meantime
-            else:
-                space.call_method(w_stream, 'flush') # XXX: ignore IOErrors?
-                    
+    def flush_all(self, space):
+        while self.streams:
+            for streamholder in self.streams.keys():
+                try:
+                    del self.streams[streamholder]
+                except KeyError:
+                    pass    # key was removed in the meantime
+                else:
+                    streamholder.autoflush(space)
+
+
+def get_autoflushher(space):
+    return space.fromcache(AutoFlusher)
+
+