backport of 3a37dfe596eb: restore pickling of dict iterators

2 files changed, 20 insertions, 24 deletions

diff --git a/pypy/objspace/std/dictmultiobject.py b/pypy/objspace/std/dictmultiobject.py
index 3cdea9563e..41235c5bc1 100644
--- a/pypy/objspace/std/dictmultiobject.py
+++ b/pypy/objspace/std/dictmultiobject.py
@@ -1432,43 +1432,27 @@ class W_BaseDictMultiIterObject(W_Root):
         At unpickling time, we just use that list
         and create an iterator on it.
         This is of course not the standard way.
-
-        XXX to do: remove this __reduce__ method and do
-        a registration with copy_reg, instead.
         """
         w_mod    = space.getbuiltinmodule('_pickle_support')
         mod      = space.interp_w(MixedModule, w_mod)
         new_inst = mod.get('dictiter_surrogate_new')
-        w_typeobj = space.type(self)
 
-        raise oefmt(space.w_TypeError,
-                    "can't pickle dictionary-keyiterator objects")
-        # XXXXXX get that working again
+        w_dict = self.iteratorimplementation.w_dict
 
-        # we cannot call __init__ since we don't have the original dict
         if isinstance(self, W_DictMultiIterKeysObject):
-            w_clone = space.allocate_instance(W_DictMultiIterKeysObject,
-                                              w_typeobj)
+            w_clone = W_DictMultiIterKeysObject(space, w_dict.iterkeys())
         elif isinstance(self, W_DictMultiIterValuesObject):
-            w_clone = space.allocate_instance(W_DictMultiIterValuesObject,
-                                              w_typeobj)
+            w_clone = W_DictMultiIterValuesObject(space, w_dict.itervalues())
         elif isinstance(self, W_DictMultiIterItemsObject):
-            w_clone = space.allocate_instance(W_DictMultiIterItemsObject,
-                                              w_typeobj)
+            w_clone = W_DictMultiIterItemsObject(space, w_dict.iteritems())
         else:
             raise oefmt(space.w_TypeError,
                         "unsupported dictiter type '%R' during pickling", self)
-        w_clone.space = space
-        w_clone.content = self.content
-        w_clone.len = self.len
-        w_clone.pos = 0
-        w_clone.setup_iterator()
+
         # spool until we have the same pos
-        while w_clone.pos < self.pos:
-            w_clone.next_entry()
-            w_clone.pos += 1
-        stuff = [w_clone.next_entry() for i in range(w_clone.pos, w_clone.len)]
-        w_res = space.newlist(stuff)
+        for x in xrange(self.iteratorimplementation.pos):
+            w_clone.descr_next(space)
+        w_res = space.call_function(space.w_list, w_clone)
         w_ret = space.newtuple([new_inst, space.newtuple([w_res])])
         return w_ret
 
diff --git a/pypy/objspace/std/test/test_dictmultiobject.py b/pypy/objspace/std/test/test_dictmultiobject.py
index 12d8654354..349222df90 100644
--- a/pypy/objspace/std/test/test_dictmultiobject.py
+++ b/pypy/objspace/std/test/test_dictmultiobject.py
@@ -1125,6 +1125,18 @@ class AppTestDictViews:
         assert (foo2, foo2_bis) in logger_copy
         assert logger_copy.issubset({(foo1, foo2_bis), (foo2, foo2_bis), (foo3, foo2_bis)})
 
+    def test_pickle(self):
+        d = {1: 1, 2: 2, 3: 3}
+        it = iter(d)
+        first = next(it)
+        reduced = it.__reduce__()
+        rebuild, args = reduced
+        new = rebuild(*args)
+        items = set(new)
+        assert len(items) == 2
+        items.add(first)
+        assert items == set(d)
+
 
 class AppTestStrategies(object):
     def setup_class(cls):