1
1
import ast
2
2
3
+
3
4
class NameAccumulator (ast .NodeVisitor ):
5
+
4
6
"""
5
7
NameAccumulator walks an AST "target" node, recursively gathering the 'id'
6
8
properties of all of the Names it finds.
7
9
"""
10
+
8
11
def __init__ (self ):
9
12
self .names = set ()
10
13
11
14
def visit_Name (self , name ):
12
15
self .names .add (name .id )
13
16
17
+
14
18
def get_module_globals (module ):
15
19
"""
16
20
Examine all of the top-level nodes in the module, and remember the names
@@ -22,13 +26,15 @@ def get_module_globals(module):
22
26
for t in node .targets :
23
27
acc .visit (t )
24
28
return acc .names
25
-
29
+
26
30
27
31
class FindFunctionAssignments (ast .NodeVisitor ):
32
+
28
33
"""
29
- Finds assignments in a function body, and accumlates the names of the
30
- assigned-to entities.
34
+ Finds assignments in a function body, and accumulates the
35
+ names of the assigned-to entities.
31
36
"""
37
+
32
38
def __init__ (self ):
33
39
self .acc = NameAccumulator ()
34
40
@@ -43,18 +49,19 @@ class FindFunctionAssignments(ast.NodeVisitor):
43
49
self .visit (func )
44
50
return self .acc .names
45
51
52
+
46
53
def insert_global_statements (module ):
47
54
"""
48
55
Finds all of the function definitions in a module, and inserts global
49
56
statements in those that assign to names that are the same as existing
50
57
module globals.
51
-
58
+
52
59
For example, insert_global_statements will transform the AST for
53
-
60
+
54
61
x = 0
55
62
def draw():
56
63
x = (x + 1) % width
57
-
64
+
58
65
into the AST for
59
66
60
67
x = 0
@@ -69,7 +76,7 @@ def insert_global_statements(module):
69
76
globals = get_module_globals (module )
70
77
needed = assigned_names .difference (args ).intersection (globals )
71
78
node .body .insert (0 , __global__ (needed ))
72
-
79
+
73
80
source = """
74
81
x = 0
75
82
@@ -85,7 +92,8 @@ def draw():
85
92
"""
86
93
module = ast .parse (source )
87
94
insert_global_statements (module )
88
- #fixed = ast.fix_missing_locations(p)
95
+ # Is this desirable?
96
+ # fixed = ast.fix_missing_locations(p)
89
97
codeobj = compile (module , __file__ , mode = 'exec' )
90
98
exec (codeobj )
91
99
0 commit comments