Commits

Armin Rigo committed b397506

Fix copy.copy() on generator objects to make sure they get their own
frame too, rather than sharing it with the original generator (which
crashes completely). Of limited use (see comments in the test).

  • Participants
  • Parent commits 656b4a4

Comments (0)

Files changed (5)

pypy/interpreter/generator.py

         new_inst = mod.get('generator_new')
         w        = space.wrap
         if self.frame:
-            w_frame = w(self.frame)
+            w_frame = self.frame._reduce_state(space)
         else:
             w_frame = space.w_None
 
             w(self.running),
             ]
 
-        return space.newtuple([new_inst, space.newtuple(tup)])
+        return space.newtuple([new_inst, space.newtuple([]),
+                               space.newtuple(tup)])
+
+    def descr__setstate__(self, space, w_args):
+        from rpython.rlib.objectmodel import instantiate
+        args_w = space.unpackiterable(w_args)
+        w_framestate, w_running = args_w
+        if space.is_w(w_framestate, space.w_None):
+            self.frame = None
+        else:
+            frame = instantiate(space.FrameClass)   # XXX fish
+            frame.descr__setstate__(space, w_framestate)
+            GeneratorIterator.__init__(self, frame)
+        self.running = self.space.is_true(w_running)
 
     def descr__iter__(self):
         """x.__iter__() <==> iter(x)"""

pypy/interpreter/pyframe.py

 from rpython.tool.stdlib_opcode import host_bytecode_spec
 
 # Define some opcodes used
-g = globals()
 for op in '''DUP_TOP POP_TOP SETUP_LOOP SETUP_EXCEPT SETUP_FINALLY
 POP_BLOCK END_FINALLY'''.split():
-    g[op] = stdlib_opcode.opmap[op]
+    globals()[op] = stdlib_opcode.opmap[op]
 HAVE_ARGUMENT = stdlib_opcode.HAVE_ARGUMENT
 
 class PyFrame(eval.Frame):
     @jit.dont_look_inside
     def descr__reduce__(self, space):
         from pypy.interpreter.mixedmodule import MixedModule
-        from pypy.module._pickle_support import maker # helper fns
         w_mod    = space.getbuiltinmodule('_pickle_support')
         mod      = space.interp_w(MixedModule, w_mod)
         new_inst = mod.get('frame_new')
-        w        = space.wrap
+        w_tup_state = self._reduce_state(space)
+        nt = space.newtuple
+        return nt([new_inst, nt([]), w_tup_state])
+
+    @jit.dont_look_inside
+    def _reduce_state(self, space):
+        from pypy.module._pickle_support import maker # helper fns
+        w = space.wrap
         nt = space.newtuple
 
         cells = self._getcells()
             w(self.instr_prev_plus_one),
             w_cells,
             ]
-
-        return nt([new_inst, nt([]), nt(tup_state)])
+        return nt(tup_state)
 
     @jit.dont_look_inside
     def descr__setstate__(self, space, w_args):

pypy/interpreter/test/test_zzpickle_and_slow.py

         pckl   = pickle.dumps(pack.mod)
         result = pickle.loads(pckl)
         assert pack.mod is result
+
+
+class AppTestGeneratorCloning:
+
+    def setup_class(cls):
+        try:
+            cls.space.appexec([], """():
+                def f(): yield 42
+                f().__reduce__()
+            """)
+        except TypeError, e:
+            if 'pickle generator' not in str(e):
+                raise
+            py.test.skip("Frames can't be __reduce__()-ed")
+
+    def test_deepcopy_generator(self):
+        import copy
+
+        def f(n):
+            for i in range(n):
+                yield 42 + i
+        g = f(4)
+        g2 = copy.deepcopy(g)
+        res = g.next()
+        assert res == 42
+        res = g2.next()
+        assert res == 42
+        g3 = copy.deepcopy(g)
+        res = g.next()
+        assert res == 43
+        res = g2.next()
+        assert res == 43
+        res = g3.next()
+        assert res == 43
+
+    def test_shallowcopy_generator(self):
+        """Note: shallow copies of generators are often confusing.
+        To start with, 'for' loops have an iterator that will not
+        be copied, and so create tons of confusion.
+        """
+        import copy
+
+        def f(n):
+            while n > 0:
+                yield 42 + n
+                n -= 1
+        g = f(2)
+        g2 = copy.copy(g)
+        res = g.next()
+        assert res == 44
+        res = g2.next()
+        assert res == 44
+        g3 = copy.copy(g)
+        res = g.next()
+        assert res == 43
+        res = g2.next()
+        assert res == 43
+        res = g3.next()
+        assert res == 43
+        g4 = copy.copy(g2)
+        for i in range(2):
+            raises(StopIteration, g.next)
+            raises(StopIteration, g2.next)
+            raises(StopIteration, g3.next)
+            raises(StopIteration, g4.next)

pypy/interpreter/typedef.py

 GeneratorIterator.typedef = TypeDef("generator",
     __repr__   = interp2app(GeneratorIterator.descr__repr__),
     __reduce__   = interp2app(GeneratorIterator.descr__reduce__),
+    __setstate__ = interp2app(GeneratorIterator.descr__setstate__),
     next       = interp2app(GeneratorIterator.descr_next,
                             descrmismatch='next'),
     send       = interp2app(GeneratorIterator.descr_send,

pypy/module/_pickle_support/maker.py

     tb = instantiate(PyTraceback)
     return space.wrap(tb)
 
-@unwrap_spec(running=int)
-def generator_new(space, w_frame, running):
-    frame = space.interp_w(PyFrame, w_frame, can_be_None=True)
-    new_generator = GeneratorIterator(frame)
-    new_generator.running = running
+def generator_new(space):
+    new_generator = instantiate(GeneratorIterator)
     return space.wrap(new_generator)
 
 @unwrap_spec(current=int, remaining=int, step=int)