[pypy-commit] pypy default: make sure to flush all _io streams when we exit the interpreter

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: 
Changeset: r52806:d4dee87e47cc
Date: 2012-02-23 12:06 +0100
http://bitbucket.org/pypy/pypy/changeset/d4dee87e47cc/

Log:	make sure to flush all _io streams when we exit the interpreter

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
@@ -28,6 +28,7 @@
         }
 
     def init(self, space):
+        MixedModule.init(self, space)
         w_UnsupportedOperation = space.call_function(
             space.w_type,
             space.wrap('UnsupportedOperation'),
@@ -35,3 +36,8 @@
             space.newdict())
         space.setattr(self, space.wrap('UnsupportedOperation'),
                       w_UnsupportedOperation)
+
+    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)
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
@@ -43,6 +43,7 @@
         self.space = space
         self.w_dict = space.newdict()
         self.__IOBase_closed = False
+        register_flushable_stream(space, self)
 
     def getdict(self, space):
         return self.w_dict
@@ -98,6 +99,8 @@
             space.call_method(self, "flush")
         finally:
             self.__IOBase_closed = True
+            flushable_streams = get_flushable_streams(space)
+            del flushable_streams[self]
 
     def flush_w(self, space):
         if self._CLOSED():
@@ -303,3 +306,32 @@
     read = interp2app(W_RawIOBase.read_w),
     readall = interp2app(W_RawIOBase.readall_w),
 )
+
+
+# ------------------------------------------------------------
+# functions to make sure that all streams are flushed on exit
+# ------------------------------------------------------------
+
+class IoState:
+    def __init__(self, space):
+        self.flushable_streams = {}
+
+def get_flushable_streams(space):
+    return space.fromcache(IoState).flushable_streams
+
+def register_flushable_stream(space, w_stream):
+    streams = get_flushable_streams(space)
+    streams[w_stream] = None
+
+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?
+                    
diff --git a/pypy/module/_io/test/test_fileio.py b/pypy/module/_io/test/test_fileio.py
--- a/pypy/module/_io/test/test_fileio.py
+++ b/pypy/module/_io/test/test_fileio.py
@@ -160,3 +160,20 @@
         f.close()
         assert repr(f) == "<_io.FileIO [closed]>"
 
+def test_flush_at_exit():
+    from pypy import conftest
+    from pypy.tool.option import make_config, make_objspace
+    from pypy.tool.udir import udir
+
+    tmpfile = udir.join('test_flush_at_exit')
+    config = make_config(conftest.option)
+    space = make_objspace(config)
+    space.appexec([space.wrap(str(tmpfile))], """(tmpfile):
+        import io
+        f = io.open(tmpfile, 'w')
+        f.write('42')
+        # no flush() and no close()
+        import sys; sys._keepalivesomewhereobscure = f
+    """)
+    space.finish()
+    assert tmpfile.read() == '42'