From 04248a8d99e518e7e833a1e2de9e3a8c1a325245 Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <solipsis@pitrou.net>
Date: Tue, 12 Oct 2010 20:51:21 +0000
Subject: Issue #3873: Speed up unpickling from file objects which have a
 peek() method.

---
 Lib/test/pickletester.py | 43 ++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 42 insertions(+), 1 deletion(-)

(limited to 'Lib')

diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py
index 33c85dcbed..7645b548a5 100644
--- a/Lib/test/pickletester.py
+++ b/Lib/test/pickletester.py
@@ -30,6 +30,21 @@ def count_opcode(code, pickle):
             n += 1
     return n
 
+
+class UnseekableIO(io.BytesIO):
+    def peek(self, *args):
+        raise NotImplementedError
+
+    def seekable(self):
+        return False
+
+    def seek(self, *args):
+        raise io.UnsupportedOperation
+
+    def tell(self):
+        raise io.UnsupportedOperation
+
+
 # We can't very well test the extension registry without putting known stuff
 # in it, but we have to be careful to restore its original state.  Code
 # should do this:
@@ -1072,9 +1087,10 @@ class AbstractPickleTests(unittest.TestCase):
         # Test the correctness of internal buffering routines when handling
         # large data.
         for proto in protocols:
-            data = (1, b'x' * (256 * 1024))
+            data = (1, min, b'xy' * (30 * 1024), len)
             dumped = self.dumps(data, proto)
             loaded = self.loads(dumped)
+            self.assertEqual(len(loaded), len(data))
             self.assertEqual(loaded, data)
 
 
@@ -1373,6 +1389,31 @@ class AbstractPicklerUnpicklerObjectTests(unittest.TestCase):
         f.seek(0)
         self.assertEqual(unpickler.load(), data2)
 
+    def _check_multiple_unpicklings(self, ioclass):
+        for proto in protocols:
+            data1 = [(x, str(x)) for x in range(2000)] + [b"abcde", len]
+            f = ioclass()
+            pickler = self.pickler_class(f, protocol=proto)
+            pickler.dump(data1)
+            pickled = f.getvalue()
+
+            N = 5
+            f = ioclass(pickled * N)
+            unpickler = self.unpickler_class(f)
+            for i in range(N):
+                if f.seekable():
+                    pos = f.tell()
+                self.assertEqual(unpickler.load(), data1)
+                if f.seekable():
+                    self.assertEqual(f.tell(), pos + len(pickled))
+            self.assertRaises(EOFError, unpickler.load)
+
+    def test_multiple_unpicklings_seekable(self):
+        self._check_multiple_unpicklings(io.BytesIO)
+
+    def test_multiple_unpicklings_unseekable(self):
+        self._check_multiple_unpicklings(UnseekableIO)
+
 
 if __name__ == "__main__":
     # Print some stuff that can be used to rewrite DATA{0,1,2}
-- 
cgit v1.2.1