Skip to content

Commit 6602852

Browse files
WaidhoferjWaidhoferj
authored
Merge pull request #3 from jupytercalpoly/03a60e3
Bifrost Tracing Update
2 parents 6bb0661 + 03a60e3 commit 6602852

File tree

2 files changed

+191
-37
lines changed

2 files changed

+191
-37
lines changed

‎bifrost_tracing/__init__.py

+7-6
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,16 @@
55
__version__ = '0.1.0'
66

77
from .bifrost_tracing import BifrostTracing, BifrostWatcher
8-
from IPython import get_ipython
8+
# from IPython import get_ipython
99

1010

1111

1212

13-
def load_ipython_extension(ipython):
14-
ipython.register_magics(BifrostTracing)
15-
vw = BifrostWatcher(ipython)
16-
ipython.events.register('post_run_cell', vw.post_run_cell)
13+
# def load_ipython_extension(ipython):
14+
# ipython.register_magics(BifrostTracing)
15+
# vw = BifrostWatcher(ipython)
16+
# ipython.events.register('post_run_cell', vw.post_run_cell)
17+
# return vw
1718

1819

19-
load_ipython_extension(get_ipython())
20+
# Watcher = load_ipython_extension(get_ipython())

‎bifrost_tracing/bifrost_tracing.py

+184-31
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,15 @@
11
"""Main module."""
22
from IPython import get_ipython
3-
from IPython.core.magic import (Magics, magics_class, cell_magic, line_magic)
3+
from IPython.core.magic import Magics, magics_class, cell_magic, line_magic
44
from IPython.core.extensions import ExtensionManager
5-
import ast
5+
from IPython import get_ipython
66

7+
import ast
78

8-
99

1010
# keep it here for reactive cells later??
1111
@magics_class
1212
class BifrostTracing(Magics):
13-
1413
@line_magic
1514
def line_tracing(self, line):
1615
return line
@@ -25,67 +24,221 @@ def __init__(self, ip):
2524
self.shell = ip
2625
self.last_x = None
2726
self.bifrost_table = {}
27+
self.plot_output = ""
28+
self.bifrost_input = ""
29+
self.visitor = None
2830

2931
def post_run_cell(self, result):
30-
ast_tree = ast.parse(result.info.raw_cell)
31-
assignVisitor = AssignVisitor()
32-
assignVisitor.visit(ast_tree)
33-
34-
for new_df in assignVisitor.new_dfs:
35-
columns = get_ipython().run_cell(new_df + '.columns').result
36-
print(columns)
37-
if new_df in self.bifrost_table:
38-
columns_set = set(columns)
39-
table_set = set(self.bifrost_table[new_df].keys())
40-
for new_col in (columns_set - table_set):
41-
self.bifrost_table[new_df][new_col] = 0
42-
else:
43-
self.bifrost_table[new_df] = {col: 0 for col in columns}
44-
32+
if not result.error_in_exec:
33+
ast_tree = ast.parse(result.info.raw_cell)
34+
35+
callVisitor = CallVisitor()
36+
callVisitor.visit(ast_tree)
37+
print(f"args={callVisitor.args}")
38+
39+
for arg in callVisitor.args:
40+
key, value = arg.split(".")
41+
if key in self.bifrost_table and value in self.bifrost_table[key]:
42+
self.bifrost_table[key][value] = self.bifrost_table[key][value] + 1
43+
elif key in self.bifrost_table and value not in self.bifrost_table[key]:
44+
self.bifrost_table[key][value] = 1
45+
else:
46+
self.bifrost_table[key] = {value: 1}
47+
# assignVisitor = AssignVisitor()
48+
# assignVisitor.visit(ast_tree)
49+
50+
# for new_df in assignVisitor.new_dfs:
51+
# columns = get_ipython().run_cell(new_df + '.columns').result
52+
# print(columns)
53+
# if new_df in self.bifrost_table:
54+
# columns_set = set(columns)
55+
# table_set = set(self.bifrost_table[new_df].keys())
56+
# for new_col in (columns_set - table_set):
57+
# self.bifrost_table[new_df][new_col] = 0
58+
# else:
59+
# self.bifrost_table[new_df] = {col: 0 for col in columns}
60+
61+
62+
class AttributeVisitor(ast.NodeVisitor):
63+
def __init__(self):
64+
self.attributes = []
4565

46-
class AttributeVisitor(ast.NodeVisitor):
4766
def visit_Attribute(self, node):
48-
self.attributes = []
49-
self.attributes.append(node.value.id + "." + node.attr)
67+
if isinstance(node, ast.Attribute):
68+
self.visit(node.value)
69+
if isinstance(node.value, ast.Name):
70+
self.attributes.append(node.attr)
71+
elif isinstance(node.value, ast.Call):
72+
callVisitor = CallVisitor()
73+
callVisitor.visit(node.value)
5074

5175

5276
class NameVisitor(ast.NodeVisitor):
53-
def visit_Name(self, node):
77+
def __init__(self):
5478
self.names = []
79+
80+
def visit_Name(self, node):
5581
self.names.append(node.id)
5682

5783

5884
class AssignVisitor(ast.NodeVisitor):
5985
def __init__(self):
6086
self.new_dfs = []
87+
self.output_var = None
88+
self.bifrost_input = None
6189

6290
def visit_Module(self, node):
6391
self.generic_visit(node)
64-
92+
6593
def visit_Assign(self, node):
6694
nameVisitor = NameVisitor()
6795
for target in node.targets:
6896
nameVisitor.visit(target)
6997
names = nameVisitor.names
7098
attributeVisitor = AttributeVisitor()
71-
attributeVisitor.visit(node.value)
99+
attributeVisitor.visit(node.value)
72100
attributes = attributeVisitor.attributes
73101
df_mask = [call == "pd.DataFrame" for call in attributes]
74-
self.new_dfs.extend([name for name, is_df in zip(names, df_mask) if is_df ])
75-
102+
plot_mask = "bifrost.plot" in attributes
103+
if "DataFrame" in attributes:
104+
self.new_dfs.extend(names)
105+
if "bifrost.plot" in attributes:
106+
self.output_var = names[-1] if len(names) else None
107+
nameVisitor = NameVisitor()
108+
nameVisitor.visit(node)
109+
self.bifrost_input = nameVisitor.names[-1]
110+
111+
112+
class SubscriptVisitor(ast.NodeVisitor):
113+
def visit_Subscript(self, node: ast.Subscript):
114+
self.subscripts = []
115+
if isinstance(node.slice, ast.Tuple):
116+
columns = [element.value for element in node.slice.elts]
117+
for column in columns:
118+
self.subscripts.append[[node.value.id, column]]
119+
elif isinstance(node.slice, ast.Compare):
120+
left = node.slice.left
121+
if isinstance(left, ast.Subscript):
122+
self.visit_Subscript(left)
123+
else:
124+
self.subscripts.append([left.value.id, left.attr])
125+
else:
126+
self.subscripts.append([node.value.id, node.slice.value])
127+
128+
129+
class CallVisitor(ast.NodeVisitor):
130+
def __init__(self):
131+
self.args = []
76132

133+
def visit_Module(self, node):
134+
self.generic_visit(node)
77135

136+
def visit_Assign(self, node):
137+
if isinstance(node.targets[0], ast.Subscript):
138+
self.get_args(node.targets[0])
139+
140+
def visit_Expr(self, node):
141+
if isinstance(node.value, ast.Call):
142+
self.visit_Call(node.value)
143+
elif isinstance(node.value, ast.Subscript):
144+
self.get_args(node.value)
145+
elif isinstance(node.value, ast.Attribute):
146+
if isinstance(node.value.value, ast.Name):
147+
self.get_args(node.value.attr, node.value.value.id)
148+
elif isinstance(node.value.value, ast.Subscript):
149+
self.get_args(node.value.value)
150+
151+
def visit_Call(self, node):
152+
attributeVisitor = AttributeVisitor()
153+
if not isinstance(node.func, ast.Attribute):
154+
return
155+
if isinstance(node.func.value, ast.Subscript):
156+
self.get_args(node.func.value)
157+
elif isinstance(node.func.value, ast.Name):
158+
func = node.func.value.id
159+
# NP case
160+
if func in ["np", "numpy"]:
161+
attributeVisitor.visit(node.func)
162+
if attributeVisitor.attributes[0] in ["mean", "std", "sum"]:
163+
args = node.args
164+
if len(args) != 0:
165+
self.get_args(args[0])
166+
else:
167+
# check keywords
168+
keywords = node.keywords
169+
for keyword in keywords:
170+
if keyword.arg == "a":
171+
self.get_args(keyword.value)
172+
# DF/PD case
173+
elif func in ["df", "pd"]:
174+
attributeVisitor.visit(node.func)
175+
attribute = attributeVisitor.attributes[0]
176+
if attribute in [
177+
"groupby",
178+
"loc",
179+
"iloc",
180+
]:
181+
args = node.args
182+
if len(args) != 0:
183+
if isinstance(node.func.value, ast.Name):
184+
self.get_args(args[0], node.func.value.id)
185+
else:
186+
self.get_args(args[0])
187+
else:
188+
# check keywords
189+
keywords = node.keywords
190+
for keyword in keywords:
191+
if attribute == "groupby" and keyword.arg == "by":
192+
self.get_args(keyword.value)
193+
194+
"""value: either ast.Subscript or ast.Attribute"""
195+
196+
def get_args(self, value, dataframe=None):
197+
# case arg is df['one']
198+
if isinstance(value, ast.Subscript):
199+
subscriptVisitor = SubscriptVisitor()
200+
subscriptVisitor.visit(value)
201+
self.args.append(
202+
f"{subscriptVisitor.subscripts[0][0]}.{subscriptVisitor.subscripts[0][1]}"
203+
)
204+
# case arg is df.one
205+
elif isinstance(value, ast.Attribute):
206+
attributeVisitor = AttributeVisitor()
207+
nameVisitor = NameVisitor()
208+
attributeVisitor.visit(value)
209+
nameVisitor.visit(value.value)
210+
df = nameVisitor.names[0]
211+
column = attributeVisitor.attributes[0]
212+
self.args.append(f"{df}.{column}")
213+
elif isinstance(value, ast.List):
214+
columns = [element.value for element in value.elts]
215+
if dataframe:
216+
for column in columns:
217+
self.args.append(f"{dataframe}.{column}")
218+
elif isinstance(value, ast.Constant):
219+
if dataframe:
220+
self.args.append(f"{dataframe}.{value.value}")
221+
222+
223+
# some_var = ...bifrost.plot()
78224
def isnotebook():
79225
try:
80226
shell = get_ipython().__class__.__name__
81-
if shell == 'ZMQInteractiveShell':
82-
return True # Jupyter notebook or qtconsole
83-
elif shell == 'TerminalInteractiveShell':
227+
if shell == "ZMQInteractiveShell":
228+
return True # Jupyter notebook or qtconsole
229+
elif shell == "TerminalInteractiveShell":
84230
return False # Terminal running IPython
85231
else:
86232
return False # Other type (?)
87233
except NameError:
88-
return False # Probably standard Python interpreter
234+
return False # Probably standard Python interpreter
235+
89236

237+
def load_ipython_extension(ipython):
238+
ipython.register_magics(BifrostTracing)
239+
vw = BifrostWatcher(ipython)
240+
ipython.events.register("post_run_cell", vw.post_run_cell)
241+
return vw
90242

91243

244+
Watcher = load_ipython_extension(get_ipython())

0 commit comments

Comments
 (0)