Some experimental support for generating NEWOBJ with proto=2, and

fixed a bug in load_newobj().

2 files changed, 70 insertions, 2 deletions

diff --git a/Lib/pickle.py b/Lib/pickle.py
index 1b364c6..e433616 100644
--- a/Lib/pickle.py
+++ b/Lib/pickle.py
@@ -77,11 +77,13 @@ class _Stop(Exception):
     def __init__(self, value):
         self.value = value
 
+# Jython has PyStringMap; it's a dict subclass with string keys
 try:
     from org.python.core import PyStringMap
 except ImportError:
     PyStringMap = None
 
+# UnicodeType may or may not be exported (normally imported from types)
 try:
     UnicodeType
 except NameError:
@@ -249,7 +251,10 @@ class Pickler:
 
         return GET + `i` + '\n'
 
-    def save(self, obj):
+    def save(self, obj,
+             _builtin_type = (int, long, float, complex, str, unicode,
+                              tuple, list, dict),
+             ):
         # Check for persistent id (defined by a subclass)
         pid = self.persistent_id(obj)
         if pid:
@@ -278,6 +283,30 @@ class Pickler:
             self.save_global(obj)
             return
 
+        # Check for instance of subclass of common built-in types
+        # XXX This block is experimental code that will go away!
+        if self.proto >= 2:
+            if isinstance(obj, _builtin_type):
+                assert t not in _builtin_type # Proper subclass
+                args = ()
+                getnewargs = getattr(obj, "__getnewargs__", None)
+                if getnewargs:
+                    args = getnewargs() # This better not reference obj
+                self.save_global(t)
+                self.save(args)
+                self.write(NEWOBJ)
+                self.memoize(obj)
+                getstate = getattr(obj, "__getstate__", None)
+                if getstate:
+                    state = getstate()
+                else:
+                    state = getattr(obj, "__dict__", None)
+                    # XXX What about __slots__?
+                if state is not None:
+                    self.save(state)
+                    self.write(BUILD)
+                return
+
         # Check copy_reg.dispatch_table
         reduce = dispatch_table.get(t)
         if reduce:
@@ -970,7 +999,7 @@ class Unpickler:
         args = self.stack.pop()
         cls = self.stack[-1]
         obj = cls.__new__(cls, *args)
-        self.stack[-1:] = obj
+        self.stack[-1] = obj
     dispatch[NEWOBJ] = load_newobj
 
     def load_global(self):
diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py
index 3d65383..ea658bb 100644
--- a/Lib/test/pickletester.py
+++ b/Lib/test/pickletester.py
@@ -300,6 +300,45 @@ class AbstractPickleTests(unittest.TestCase):
                 y = self.loads(s)
                 self.assert_(x is y, (proto, x, s, y))
 
+    def test_newobj_tuple(self):
+        x = MyTuple([1, 2, 3], foo=42, bar="hello")
+        s = self.dumps(x, 2)
+        y = self.loads(s)
+        self.assertEqual(tuple(x), tuple(y))
+        self.assertEqual(x.__dict__, y.__dict__)
+
+    def test_newobj_list(self):
+        x = MyList([1, 2, 3], foo=42, bar="hello")
+        s = self.dumps(x, 2)
+        y = self.loads(s)
+        self.assertEqual(list(x), list(y))
+        self.assertEqual(x.__dict__, y.__dict__)
+
+class MyTuple(tuple):
+    def __new__(cls, *args, **kwds):
+        # Ignore **kwds
+        return tuple.__new__(cls, *args)
+    def __getnewargs__(self):
+        return (tuple(self),)
+    def __init__(self, *args, **kwds):
+        for k, v in kwds.items():
+            setattr(self, k, v)
+
+class MyList(list):
+    def __new__(cls, *args, **kwds):
+        # Ignore **kwds
+        return list.__new__(cls, *args)
+    def __init__(self, *args, **kwds):
+        for k, v in kwds.items():
+            setattr(self, k, v)
+    def __getstate__(self):
+        return list(self), self.__dict__
+    def __setstate__(self, arg):
+        lst, dct = arg
+        for x in lst:
+            self.append(x)
+        self.__init__(**dct)
+
 class AbstractPickleModuleTests(unittest.TestCase):
 
     def test_dump_closed_file(self):