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

"""

Here we add a second version of every function defined in a module.
The new version returns the local bindings at the time return is called
in the original function, along with the original return value,
under the name __return_value__.

This new version of the function can be accessed under the name
    myOriginalFunction_exit_context
and will preserve the same arguments, defaults, etc
"""

import compiler,copy
from compiler.ast import Return,Assign,Name,AssName,CallFunc,Function

returnLocals = Return(CallFunc(Name('locals'), [], None, None))

class AddExitContextVisitor:
    
    def visitFunction(self, node, *args):
        node.name = node.name+'_exit_context'
        for child in node.getChildNodes():
            self.visit(child,*args)
        
        if not node.code.hasReturns:
            node.code.nodes.append(returnLocals)
    
    def visitStmt(self, node, *args):
        insertions = {}
        for i,child in enumerate(node.nodes):
            if isinstance(child, Return):
                insertions[i] = [Assign([AssName('__return_value__', 'OP_ASSIGN')], child.value),
                                 returnLocals]
            
            elif isinstance(child, Function):
                insertions[i] = [copy.deepcopy(child), child]
                            
            self.visit(child,*args)
        
        for index,insertion in insertions.iteritems():
            node.nodes[index:index+1] = insertion
    
        node.hasReturns = bool(insertions)
        

def addExitContextFunctions(source):
    ast = compiler.parse(source)
    walker = compiler.visitor.ASTVisitor()
    walker.preorder(ast, AddExitContextVisitor())
    compiler.misc.set_filename('exit_contextified_'+__file__, ast)
    return compiler.pycodegen.ModuleCodeGenerator(ast).getCode()


    
if __name__ == '__main__':
    # testing
    exec addExitContextFunctions("""

def placeholder():
    pass # hmm! bug!

def function(a,b,c):
    d = a+b
    return d,c

def something_else(r,e=True,t=False):
    l = [1,2,3]
    if r > 5:
        pig = 32
        l.append(e)
    else:
        l.append(t)
    return l

def noReturn(a,b,c,e):
    l = (a,b)
    m = (c,e)
    n = (c,c,c,e)
    o = (e,e,a)
    p = (e,c,l,m,n,o)

""")
    
    for src in [ 
        "function_exit_context('cat','dog','aardvark')",
        "something_else_exit_context(3, 'alpha', 'beta')",
        "something_else(12, 'here', 'there')",
        "something_else_exit_context(12, 'here', 'there')",
        "noReturn(1,2,3,4)",
        "noReturn_exit_context(1,2,3,4)"
    ]:
        print '\n',src,'=>\n',eval(src)