diff --git a/Lib/pickle.py b/Lib/pickle.py index 702b0b3..6dee0d3 100644 --- a/Lib/pickle.py +++ b/Lib/pickle.py @@ -702,7 +702,28 @@ class _Pickler: self.memoize(obj) dispatch[bytes] = save_bytes + def save_oldstr(self, obj): + if self.bin: + n = len(obj) + if n < 256: + self.write(SHORT_BINSTRING + pack("', '<\\\u1234>', '<\n>', '<\\>', '<\\\U00012345>', @@ -1840,6 +1859,8 @@ class AbstractPickleTests(unittest.TestCase): self.assertIn(b'\nL64206', s) # LONG elif proto < 2: self.assertIn(b'M\xce\xfa', s) # BININT2 + elif proto < 3: + self.assertIn(b'U\x04FACE', s) # SHORT_BINSTRING elif proto < 4: self.assertIn(b'X\x04\x00\x00\x00FACE', s) # BINUNICODE else: @@ -1860,6 +1881,8 @@ class AbstractPickleTests(unittest.TestCase): self.assertIn(b'\nL64206', s) # LONG elif proto < 2: self.assertIn(b'M\xce\xfa', s) # BININT2 + elif proto < 3: + self.assertIn(b'U\x04FACE', s) # SHORT_BINSTRING elif proto < 4: self.assertIn(b'X\x04\x00\x00\x00FACE', s) # BINUNICODE else: diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py index e6c5d08..c5950f4 100644 --- a/Lib/test/test_pickle.py +++ b/Lib/test/test_pickle.py @@ -49,9 +49,9 @@ class PyPicklerTests(AbstractPickleTests): pickler = pickle._Pickler unpickler = pickle._Unpickler - def dumps(self, arg, proto=None): + def dumps(self, arg, proto=None, **kwargs): f = io.BytesIO() - p = self.pickler(f, proto) + p = self.pickler(f, proto, **kwargs) p.dump(arg) f.seek(0) return bytes(f.read()) @@ -72,8 +72,8 @@ class InMemoryPickleTests(AbstractPickleTests, AbstractUnpickleTests, AttributeError, ValueError, struct.error, IndexError, ImportError) - def dumps(self, arg, protocol=None): - return pickle.dumps(arg, protocol) + def dumps(self, arg, protocol=None, **kwargs): + return pickle.dumps(arg, protocol, **kwargs) def loads(self, buf, **kwds): return pickle.loads(buf, **kwds) @@ -81,12 +81,12 @@ class InMemoryPickleTests(AbstractPickleTests, AbstractUnpickleTests, class PersistentPicklerUnpicklerMixin(object): - def dumps(self, arg, proto=None): + def dumps(self, arg, proto=None, **kwargs): class PersPickler(self.pickler): def persistent_id(subself, obj): return self.persistent_id(obj) f = io.BytesIO() - p = PersPickler(f, proto) + p = PersPickler(f, proto, **kwargs) p.dump(arg) return f.getvalue() diff --git a/Modules/_pickle.c b/Modules/_pickle.c index e65b88e..201c345 100644 --- a/Modules/_pickle.c +++ b/Modules/_pickle.c @@ -2237,9 +2237,76 @@ write_unicode_binary(PicklerObject *self, PyObject *obj) } static int +save_oldstr(PicklerObject *self, PyObject *obj) +{ + Py_ssize_t size; + const char *data; + PyObject *repr = NULL; + + assert(PyUnicode_Check(obj)); + assert(PyUnicode_IS_ASCII(obj)); + if (!self->bin) { + const char string_op = STRING; + + if (!(repr = PyObject_ASCII(obj))) + return -1; + assert(PyUnicode_Check(repr)); + assert(PyUnicode_IS_ASCII(repr)); + + data = PyUnicode_AsUTF8AndSize(repr, &size); + if (data == NULL) + goto err; + if (_Pickler_Write(self, &string_op, 1) < 0) + goto err; + if (_Pickler_Write(self, data, size) < 0) + goto err; + if (_Pickler_Write(self, "\n", 1) < 0) + goto err; + + Py_DECREF(repr); + } + else { + Py_ssize_t len; + char header[5]; + + data = PyUnicode_AsUTF8AndSize(obj, &size); + if (size < 256) { + header[0] = SHORT_BINSTRING; + header[1] = size; + len = 2; + } + else { + assert(size <= 0x7fffffffL); + header[0] = BINSTRING; + header[1] = (unsigned char)(size & 0xff); + header[2] = (unsigned char)((size >> 8) & 0xff); + header[3] = (unsigned char)((size >> 16) & 0xff); + header[4] = (unsigned char)((size >> 24) & 0xff); + len = 5; + } + + if (_Pickler_Write(self, header, len) < 0) + return -1; + if (_Pickler_Write(self, data, size) < 0) + return -1; + } + + return 0; + + err: + Py_XDECREF(repr); + return -1; +} + +static int save_unicode(PicklerObject *self, PyObject *obj) { - if (self->bin) { + if (self->fix_imports && self->proto < 3 && + PyUnicode_IS_ASCII(obj) && PyUnicode_GET_LENGTH(obj) <= 0x7fffffffL) { + if (save_oldstr(self, obj) < 0) + return -1; + } + else if (self->bin) { if (write_unicode_binary(self, obj) < 0) return -1; }