[pypy-commit] pypy py3.5: Pickle protocol 4: Call __getnewargs_ex__(), implemented when __new__() requires keyword arguments.

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: py3.5
Changeset: r88056:82efc04d5dda
Date: 2016-11-02 09:04 +0100
http://bitbucket.org/pypy/pypy/changeset/82efc04d5dda/

Log:	Pickle protocol 4: Call __getnewargs_ex__(), implemented when
	__new__() requires keyword arguments.

diff --git a/pypy/objspace/std/objectobject.py b/pypy/objspace/std/objectobject.py
--- a/pypy/objspace/std/objectobject.py
+++ b/pypy/objspace/std/objectobject.py
@@ -18,17 +18,27 @@
     import copyreg
     return copyreg._reduce_ex(obj, proto)
 
-def reduce_2(obj):
-    cls = obj.__class__
+def _getnewargs(obj):
 
     try:
-        getnewargs = obj.__getnewargs__
+        getnewargs = obj.__getnewargs_ex__
     except AttributeError:
-        args = ()
+        try:
+            getnewargs = obj.__getnewargs__
+        except AttributeError:
+            args = ()
+        else:
+            args = getnewargs()
+        kwargs = None
     else:
-        args = getnewargs()
-        if not isinstance(args, tuple):
-            raise TypeError("__getnewargs__ should return a tuple")
+        args, kwargs = getnewargs()
+
+    if not isinstance(args, tuple):
+        raise TypeError("__getnewargs__ should return a tuple")
+    return args, kwargs
+
+def _getstate(obj):
+    cls = obj.__class__
 
     try:
         getstate = obj.__getstate__
@@ -48,15 +58,32 @@
                 state = state, slots
     else:
         state = getstate()
+    return state
 
+def reduce_2(obj, proto):
+    cls = obj.__class__
+
+    import copyreg
+
+    args, kwargs = _getnewargs(obj)
+
+    if not kwargs:
+       newobj = copyreg.__newobj__
+       args2 = (cls,) + args
+    elif proto >= 4:
+       newobj = copyreg.__newobj_ex__
+       args2 = (cls, args, kwargs)
+    else:
+       raise ValueError("must use protocol 4 or greater to copy this "
+                        "object; since __getnewargs_ex__ returned "
+                        "keyword arguments.")
+
+    state = _getstate(obj)
     listitems = iter(obj) if isinstance(obj, list) else None
     dictitems = iter(obj.items()) if isinstance(obj, dict) else None
 
-    import copyreg
-    newobj = copyreg.__newobj__
+    return newobj, args2, state, listitems, dictitems
 
-    args2 = (cls,) + args
-    return newobj, args2, state, listitems, dictitems
 
 def slotnames(cls):
     if not isinstance(cls, type):
@@ -169,9 +196,9 @@
 
 @unwrap_spec(proto=int)
 def descr__reduce__(space, w_obj, proto=0):
+    w_proto = space.wrap(proto)
     if proto >= 2:
-        return reduce_2(space, w_obj)
-    w_proto = space.wrap(proto)
+        return reduce_2(space, w_obj, w_proto)
     return reduce_1(space, w_obj, w_proto)
 
 @unwrap_spec(proto=int)
diff --git a/pypy/objspace/std/test/test_obj.py b/pypy/objspace/std/test/test_obj.py
--- a/pypy/objspace/std/test/test_obj.py
+++ b/pypy/objspace/std/test/test_obj.py
@@ -59,6 +59,22 @@
         s = X().__reduce__()
         assert s[-1] == ':-)'
 
+    def test_getnewargs_ex(self):
+        class NamedInt(int):
+            def __new__(cls, name, **kwargs):
+                if len(kwargs) == 0:
+                    raise TypeError("name and value must be specified")
+                self = int.__new__(cls, kwargs['value'])
+                self._name = name
+                return self
+            def __getnewargs_ex__(self):
+                return (self._name,), dict(value=int(self))
+        import copyreg
+        assert NamedInt("Name", value=42).__reduce__(4) == (
+            copyreg.__newobj_ex__,
+            (NamedInt, ('Name',), dict(value=42)),
+            dict(_name='Name'), None, None)
+
     def test_default_format(self):
         class x(object):
             def __str__(self):