Issue #22003: When initialized from a bytes object, io.BytesIO() now

pitrou · pitrou · commit cc66a73d27ba · 2014-07-29T19:41:11.000-04:00
defers making a copy until it is mutated, improving performance and
memory use on some use cases.

Patch by David Wilson.
diff --git a/Lib/test/test_memoryio.py b/Lib/test/test_memoryio.py
@@ -9,6 +9,7 @@
 import io
 import _pyio as pyio
 import pickle
+import sys
 
 class MemorySeekTestMixin:
 
@@ -711,12 +712,57 @@ def test_setstate(self):
 
     @support.cpython_only
     def test_sizeof(self):
-        basesize = support.calcobjsize('P2nN2Pn')
+        basesize = support.calcobjsize('P2nN2PnP')
         check = self.check_sizeof
         self.assertEqual(object.__sizeof__(io.BytesIO()), basesize)
         check(io.BytesIO(), basesize )
-        check(io.BytesIO(b'a'), basesize + 1 + 1 )
-        check(io.BytesIO(b'a' * 1000), basesize + 1000 + 1 )
+        check(io.BytesIO(b'a'), basesize + 1 )
+        check(io.BytesIO(b'a' * 1000), basesize + 1000)
+
+    # Various tests of copy-on-write behaviour for BytesIO.
+
+    def _test_cow_mutation(self, mutation):
+        # Common code for all BytesIO copy-on-write mutation tests.
+        imm = b' ' * 1024
+        old_rc = sys.getrefcount(imm)
+        memio = self.ioclass(imm)
+        self.assertEqual(sys.getrefcount(imm), old_rc + 1)
+        mutation(memio)
+        self.assertEqual(sys.getrefcount(imm), old_rc)
+
+    @support.cpython_only
+    def test_cow_truncate(self):
+        # Ensure truncate causes a copy.
+        def mutation(memio):
+            memio.truncate(1)
+        self._test_cow_mutation(mutation)
+
+    @support.cpython_only
+    def test_cow_write(self):
+        # Ensure write that would not cause a resize still results in a copy.
+        def mutation(memio):
+            memio.seek(0)
+            memio.write(b'foo')
+        self._test_cow_mutation(mutation)
+
+    @support.cpython_only
+    def test_cow_setstate(self):
+        # __setstate__ should cause buffer to be released.
+        memio = self.ioclass(b'foooooo')
+        state = memio.__getstate__()
+        def mutation(memio):
+            memio.__setstate__(state)
+        self._test_cow_mutation(mutation)
+
+    @support.cpython_only
+    def test_cow_mutable(self):
+        # BytesIO should accept only Bytes for copy-on-write sharing, since
+        # arbitrary buffer-exporting objects like bytearray() aren't guaranteed
+        # to be immutable.
+        ba = bytearray(1024)
+        old_rc = sys.getrefcount(ba)
+        memio = self.ioclass(ba)
+        self.assertEqual(sys.getrefcount(ba), old_rc)
 
 class CStringIOTest(PyStringIOTest):
     ioclass = io.StringIO
diff --git a/Misc/ACKS b/Misc/ACKS
@@ -1456,6 +1456,7 @@ Sue Williams
 Carol Willing
 Steven Willis
 Frank Willison
+David Wilson
 Geoff Wilson
 Greg V. Wilson
 J Derek Wilson
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -113,6 +113,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #22003: When initialized from a bytes object, io.BytesIO() now
+  defers making a copy until it is mutated, improving performance and
+  memory use on some use cases.  Patch by David Wilson.
+
 - Issue #22018: On Windows, signal.set_wakeup_fd() now also supports sockets.
   A side effect is that Python depends to the WinSock library.
 
diff --git a/Modules/_io/bytesio.c b/Modules/_io/bytesio.c
