1717import ast
1818
1919
20+ # For details on the AST, see https://docs.python.org/3/library/ast.html
21+ # and https://greentreesnakes.readthedocs.io/en/latest/nodes.html
22+
23+ # This class is used to instrument a Python script with calls to a
24+ # RunningScript instance. The RunningScript instance is used to
25+ # track the execution of the script, and can be used to pause and
26+ # resume the script. We inherit from ast.NodeTransformer, which
27+ # allows us to modify the AST of the script. We override the visit
28+ # method for each type of node that we want to instrument.
2029class ScriptInstrumentor (ast .NodeTransformer ):
2130 pre_line_instrumentation = """
2231RunningScript.instance.pre_line_instrumentation('{}', {}, globals(), locals())
@@ -38,93 +47,162 @@ def __init__(self, filename):
3847 self .filename = filename
3948 self .in_try = False
4049
41- # These are statements which should have an enter and leave
42- # (In retrospect, this isn't always true, eg, for 'if')
43- def track_enter_leave_lineno (self , node ):
50+ # What we're trying to do is wrap executable statements in a while True try/except block
51+ # For example if the input code is "print('HI')", we want to transform it to:
52+ # while True:
53+ # try:
54+ # RunningScript.instance.pre_line_instrumentation('myfile.py', 1, globals(), locals())
55+ # --> print('HI') <-- This is the original node
56+ # break
57+ # except:
58+ # retry_needed = RunningScript.instance.exception_instrumentation('myfile.py', 1)
59+ # if retry_needed:
60+ # continue
61+ # else:
62+ # break
63+ # finally:
64+ # RunningScript.instance.post_line_instrumentation('myfile.py', 1)
65+ # This allows us to retry statements that raise exceptions
66+ def track_enter_leave (self , node ):
67+ # Determine if we're in a try block
4468 in_try = self .in_try
4569 if not in_try and type (node ) in (ast .Try , ast .TryStar ):
4670 self .in_try = True
71+ # Visit the children of the node
4772 node = self .generic_visit (node )
4873 if not in_try and type (node ) in (ast .Try , ast .TryStar ):
4974 self .in_try = False
50- enter = ast .parse (
75+ # ast.parse returns a module, so we need to extract
76+ # the first element of the body which is the node
77+ pre_line = ast .parse (
5178 self .pre_line_instrumentation .format (self .filename , node .lineno )
5279 ).body [0 ]
53- leave = ast .parse (
80+ post_line = ast .parse (
5481 self .post_line_instrumentation .format (self .filename , node .lineno )
5582 ).body [0 ]
5683 true_node = ast .Constant (True )
5784 break_node = ast .Break ()
58- for new_node in (enter , leave , true_node , break_node ):
85+ for new_node in (pre_line , post_line , true_node , break_node ):
86+ # Copy source location from the original node to our new nodes
5987 ast .copy_location (new_node , node )
6088
61- # This is the code for "if 1: ..."
62- inhandler = ast .parse (
89+ # Create the exception handler code node. This results in multiple nodes
90+ # because we have a top level assignment and if statement
91+ exception_handler = ast .parse (
6392 self .exception_instrumentation .format (self .filename , node .lineno )
6493 ).body
65- for new_node in inhandler :
94+ for new_node in exception_handler :
6695 ast .copy_location (new_node , node )
96+ # Recursively yield the children of the new_node and copy in source locations
97+ # It's actually surprising how many nodes are nested in the new_node
6798 for new_node2 in ast .walk (new_node ):
6899 ast .copy_location (new_node2 , node )
69- excepthandler = ast .ExceptHandler (expr = None , name = None , body = inhandler )
100+ # Create an exception handler node to wrap the exception handler code
101+ excepthandler = ast .ExceptHandler (type = None , name = None , body = exception_handler )
70102 ast .copy_location (excepthandler , node )
103+ # If we're not already in a try block, we need to wrap the node in a while loop
71104 if not self .in_try :
72105 try_node = ast .Try (
73- body = [enter , node , break_node ],
106+ # pre_line is the pre_line_instrumentation, node is the original node
107+ # and if the code is executed without an exception, we break
108+ body = [pre_line , node , break_node ],
109+ # Pass in the handler we created above
74110 handlers = [excepthandler ],
111+ # No else block
75112 orelse = [],
76- finalbody = [leave ],
113+ # The try / except finally block is the post_line_instrumentation
114+ finalbody = [post_line ],
77115 )
78116 ast .copy_location (try_node , node )
79117 while_node = ast .While (test = true_node , body = [try_node ], orelse = [])
80118 ast .copy_location (while_node , node )
81119 return while_node
120+ # We're already in a try block, so we just need to wrap the node in a try block
82121 else :
83122 try_node = ast .Try (
84- body = [enter , node ],
123+ body = [pre_line , node ],
85124 handlers = [],
86125 orelse = [],
87- finalbody = [leave ],
126+ finalbody = [post_line ],
88127 )
89128 ast .copy_location (try_node , node )
90129 return try_node
91130
92- visit_Assign = track_enter_leave_lineno
93- visit_AugAssign = track_enter_leave_lineno
94- visit_Delete = track_enter_leave_lineno
95- visit_Print = track_enter_leave_lineno
96- visit_Assert = track_enter_leave_lineno
97- visit_Import = track_enter_leave_lineno
98- visit_ImportFrom = track_enter_leave_lineno
99- visit_Exec = track_enter_leave_lineno
100- # Global
101- visit_Expr = track_enter_leave_lineno
102-
103- # These statements can be reached, but they change
104- # control flow and are never exited.
105- def track_reached_lineno (self , node ):
131+ # Call the pre_line_instrumentation ONLY and then exceute the node
132+ def track_reached (self , node ):
133+ # Determine if we're in a try block, this is used by track_enter_leave
134+ in_try = self .in_try
135+ if not in_try and type (node ) in (ast .Try , ast .TryStar ):
136+ self .in_try = True
137+
138+ # Visit the children of the node
106139 node = self .generic_visit (node )
107- reach = ast .parse (
140+ pre_line = ast .parse (
108141 self .pre_line_instrumentation .format (self .filename , node .lineno )
109142 ).body [0 ]
110- ast .copy_location (reach , node )
143+ ast .copy_location (pre_line , node )
111144
112- n = ast .Num (n = 1 )
145+ # Create a simple constant node with the value 1 that we can use with our If node
146+ n = ast .Constant (value = 1 )
113147 ast .copy_location (n , node )
114- if_node = ast .If (test = n , body = [reach , node ], orelse = [])
148+ # The if_node is effectively a noop that holds the preline & node that we need to execute
149+ if_node = ast .If (test = n , body = [pre_line , node ], orelse = [])
115150 ast .copy_location (if_node , node )
116151 return if_node
117152
118- visit_With = track_reached_lineno
119- visit_FunctionDef = track_reached_lineno
120- visit_ClassDef = track_reached_lineno
121- visit_For = track_reached_lineno
122- visit_While = track_reached_lineno
123- visit_If = track_reached_lineno
124- visit_Try = track_reached_lineno
125- visit_TryStar = track_reached_lineno
126- visit_Pass = track_reached_lineno
127- visit_Return = track_reached_lineno
128- visit_Raise = track_enter_leave_lineno
129- visit_Break = track_reached_lineno
130- visit_Continue = track_reached_lineno
153+ # Notes organized (including newlines) per https://docs.python.org/3/library/ast.html#abstract-grammar
154+ # Nodes that change control flow are processed by track_reached, otherwise we track_enter_leave
155+ visit_FunctionDef = track_reached
156+ visit_AsyncFunctionDef = track_reached
157+
158+ visit_ClassDef = track_reached
159+ visit_Return = track_reached
160+
161+ visit_Delete = track_enter_leave
162+ visit_Assign = track_enter_leave
163+ visit_TypeAlias = track_enter_leave
164+ visit_AugAssign = track_enter_leave
165+ visit_AnnAssign = track_enter_leave
166+
167+ visit_For = track_reached
168+ visit_AsyncFor = track_reached
169+ visit_While = track_reached
170+ visit_If = track_reached
171+ visit_With = track_reached
172+ visit_AsyncWith = track_reached
173+
174+ # We can track the match statement but not any of the case statements
175+ # because they must come unaltered after the match statement
176+ visit_Match = track_reached
177+
178+ visit_Raise = track_enter_leave
179+ visit_Try = track_reached
180+ visit_TryStar = track_reached
181+ visit_Assert = track_enter_leave
182+
183+ visit_Import = track_enter_leave
184+ visit_ImportFrom = track_enter_leave
185+
186+ visit_Global = track_enter_leave
187+ visit_Nonlocal = track_enter_leave
188+ visit_Expr = track_enter_leave
189+ visit_Pass = track_reached
190+ visit_Break = track_reached
191+ visit_Continue = track_reached
192+
193+ # expr nodes: mostly subnodes in assignments or return statements
194+ # TODO: Should we handle the following:
195+ # visit_NamedExpr = track_enter_leave
196+ # visit_Lambda = track_enter_leave
197+ # visit_IfExp = track_enter_leave
198+ # visit_Await = track_reached
199+ # visit_Yield = track_reached
200+ # visit_YieldFrom = track_reached
201+ # visit_Call = track_reached
202+ # visit_JoinedStr = track_enter_leave
203+ # visit_Constant = track_enter_leave
204+
205+ # All the expr_context, boolop, operator, unaryop, cmpop nodes are not modified
206+ # ExceptHandler must follow try or tryStar so don't modify it
207+ # Can't modify any of pattern nodes (case) because they have to come unaltered after match
208+ # Ignore the type_ignore and type_param nodes
0 commit comments