Welcome, guest | Sign In | My Account | Store | Cart

Copying Generators (Python recipe) by kay schluehr
ActiveState Code (http://code.activestate.com/recipes/528949/)

This recipe presents copy_generator(...) which a pure Python function keeping a running generator object and returns a copy of the generator object being in the same state as the original generator object.

Python, 124 lines
Download
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
###
#
#  W A R N I N G 
#  
#  This recipe is obsolete! 
#  
#  When you are looking for copying and pickling functionality for generators
#  implemented in pure Python download the 
#
#             generator_tools 
#
#  package at the cheeseshop or at www.fiber-space.de
#              
###


import new
import copy
import types
import sys
from opcode import*


def copy_generator(f_gen):
    '''
    Function used to copy a generator object.

    @param f_gen: generator object.
    @return: pair (g_gen, g) where g_gen is a new generator object and g a generator
             function g producing g_gen. The function g is created from f_gen.gi_frame.

    Usage: function copies a running generator.

        def inc(start, step = 1):
            i = start
            while True:
                yield i
                i+= step

        >>> inc_gen = inc(3)
        >>> inc_gen.next()
        3
        >>> inc_gen.next()
        4
        >>> inc_gen_c, inc_c = copy_generator(inc_gen)
        >>> inc_gen_c.next() == inc_gen.next()
        True
        >>> inc_gen_c.next()
        6

    Implementation strategy:

        Inspecting the frame of a running generator object f provides following important
        information about the state of the generator:

           - the values of bound locals inside the generator object
           - the last bytecode being executed

        This state information of f is restored in a new function generator g in the following way:

           - the signature of g is defined by the locals of f ( co_varnames of f ). So we can pass the
             locals to g inspected from the current frame of running f. Yet unbound locals are assigned
             to None.

             All locals will be deepcopied. If one of the locals is a generator object it will be copied
             using copy_generator. If a local is not copyable it will be assigned directly. Shared state
             is therefore possible.

           - bytecode hack. A JUMP_ABSOLUTE bytecode instruction is prepended to the bytecode of f with
             an offset pointing to the next unevaluated bytecode instruction of f.

    Corner cases:

        - an unstarted generator ( last instruction = -1 ) will be just cloned.

        - if a generator has been already closed ( gi_frame = None ) a ValueError exception
          is raised.

    '''
    if not f_gen.gi_frame:
        raise ValueError("Can't copy closed generator")
    f_code = f_gen.gi_frame.f_code
    offset = f_gen.gi_frame.f_lasti
    locals = f_gen.gi_frame.f_locals

    if offset == -1:  # clone the generator
        argcount = f_code.co_argcount
    else:
        # bytecode hack - insert jump to current offset 
        # the offset depends on the version of the Python interpreter
        if sys.version_info[:2] == (2,4):
            offset +=4
        elif sys.version_info[:2] == (2,5):
            offset +=5
        start_sequence = (opmap["JUMP_ABSOLUTE"],)+divmod(offset, 256)[::-1]
        modified_code = "".join([chr(op) for op in start_sequence])+f_code.co_code
        argcount = f_code.co_nlocals

    varnames = list(f_code.co_varnames)
    for i, name in enumerate(varnames):
        loc = locals.get(name)
        if isinstance(loc, types.GeneratorType):
            varnames[i] = copy_generator(loc)[0]
        else:
            try:
                varnames[i] = copy.deepcopy(loc)
            except TypeError:
                varnames[i] = loc

    new_code = new.code(argcount,
             f_code.co_nlocals,
             f_code.co_stacksize,
             f_code.co_flags,
             modified_code,
             f_code.co_consts,
             f_code.co_names,
             f_code.co_varnames,
             f_code.co_filename,
             f_code.co_name,
             f_code.co_firstlineno,
             f_code.co_lnotab)
    g = new.function(new_code, globals(),)
    g_gen = g(*varnames)
    return g_gen, g

Unless you are using Stackless Python but rely on CPythons standard library instead there was yet no way to copy generators. The algorithm is surprisingly simple but relies also on a bytecode hack which doesn't make the recipe portable across Python implementations.

The current implementation is tested with Python version 2.4 and 2.5.

Tags: algorithms

3 comments

Klaus Muller 16 years, 7 months ago  # | flag

Error found. Kay, this and the sister recipe for pickling generators are great and badly needed in Python!

That throws an error "XXX lineno: 5, opcode: 0". If I replace the "for" with a "while", it works beautifully.

Klaus Muller 16 years, 7 months ago  # | flag

Error found. Kay, this and the sister recipe for pickling generators are great and badly needed in Python! I ran an example in Python 2.4.4:

import copyinggenerators as cp
def inc(par):
    loc=par
    for k in range(5):
        yield loc
        loc+=1

i=inc(10)
i.next()
ci,cici=cp.copy_generator(i)
print ci.next()==i.next()

That throws an error "XXX lineno: 5, opcode: 0".
If I replace the "for" with a "while", it works beautifully.
Dima Tisnek 12 years ago  # | flag

Awesome!