viztracer 1.1.0__cp314-cp314-win32.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of viztracer might be problematic. Click here for more details.

Files changed (109) hide show
  1. viztracer/__init__.py +19 -0
  2. viztracer/__main__.py +8 -0
  3. viztracer/attach.py +67 -0
  4. viztracer/attach_process/LICENSE +203 -0
  5. viztracer/attach_process/__init__.py +0 -0
  6. viztracer/attach_process/add_code_to_python_process.py +582 -0
  7. viztracer/attach_process/attach_x86.dll +0 -0
  8. viztracer/attach_process/inject_dll_amd64.exe +0 -0
  9. viztracer/attach_process/linux_and_mac/lldb_prepare.py +54 -0
  10. viztracer/attach_process/run_code_on_dllmain_amd64.dll +0 -0
  11. viztracer/attach_process/run_code_on_dllmain_x86.dll +0 -0
  12. viztracer/cellmagic.py +70 -0
  13. viztracer/code_monkey.py +353 -0
  14. viztracer/decorator.py +164 -0
  15. viztracer/event_base.py +81 -0
  16. viztracer/functree.py +135 -0
  17. viztracer/html/flamegraph.html +34 -0
  18. viztracer/html/trace_viewer_embedder.html +203 -0
  19. viztracer/html/trace_viewer_full.html +10207 -0
  20. viztracer/main.py +699 -0
  21. viztracer/modules/eventnode.c +172 -0
  22. viztracer/modules/eventnode.h +73 -0
  23. viztracer/modules/pythoncapi_compat.h +1726 -0
  24. viztracer/modules/quicktime.c +177 -0
  25. viztracer/modules/quicktime.h +104 -0
  26. viztracer/modules/snaptrace.c +2205 -0
  27. viztracer/modules/snaptrace.h +134 -0
  28. viztracer/modules/snaptrace_member.c +483 -0
  29. viztracer/modules/util.c +45 -0
  30. viztracer/modules/util.h +22 -0
  31. viztracer/modules/vcompressor/vc_dump.c +1131 -0
  32. viztracer/modules/vcompressor/vc_dump.h +49 -0
  33. viztracer/modules/vcompressor/vcompressor.c +396 -0
  34. viztracer/modules/vcompressor/vcompressor.h +15 -0
  35. viztracer/patch.py +307 -0
  36. viztracer/report_builder.py +311 -0
  37. viztracer/snaptrace.cp314-win32.pyd +0 -0
  38. viztracer/snaptrace.pyi +77 -0
  39. viztracer/util.py +196 -0
  40. viztracer/vcompressor.cp314-win32.pyd +0 -0
  41. viztracer/vcompressor.pyi +10 -0
  42. viztracer/viewer.py +528 -0
  43. viztracer/vizcounter.py +20 -0
  44. viztracer/vizevent.py +31 -0
  45. viztracer/vizlogging.py +20 -0
  46. viztracer/vizobject.py +28 -0
  47. viztracer/vizplugin.py +143 -0
  48. viztracer/viztracer.py +472 -0
  49. viztracer/web_dist/LICENSE +189 -0
  50. viztracer/web_dist/index.html +127 -0
  51. viztracer/web_dist/service_worker.js +279 -0
  52. viztracer/web_dist/trace_processor +300 -0
  53. viztracer/web_dist/v52.0-6b9586def/assets/MaterialSymbolsOutlined.woff2 +0 -0
  54. viztracer/web_dist/v52.0-6b9586def/assets/Roboto-100.woff2 +0 -0
  55. viztracer/web_dist/v52.0-6b9586def/assets/Roboto-300.woff2 +0 -0
  56. viztracer/web_dist/v52.0-6b9586def/assets/Roboto-400.woff2 +0 -0
  57. viztracer/web_dist/v52.0-6b9586def/assets/Roboto-500.woff2 +0 -0
  58. viztracer/web_dist/v52.0-6b9586def/assets/RobotoCondensed-Light.woff2 +0 -0
  59. viztracer/web_dist/v52.0-6b9586def/assets/RobotoCondensed-Regular.woff2 +0 -0
  60. viztracer/web_dist/v52.0-6b9586def/assets/RobotoMono-Regular.woff2 +0 -0
  61. viztracer/web_dist/v52.0-6b9586def/assets/brand.png +0 -0
  62. viztracer/web_dist/v52.0-6b9586def/assets/catapult_trace_viewer.html +3946 -0
  63. viztracer/web_dist/v52.0-6b9586def/assets/catapult_trace_viewer.js +7539 -0
  64. viztracer/web_dist/v52.0-6b9586def/assets/favicon.png +0 -0
  65. viztracer/web_dist/v52.0-6b9586def/assets/logo-128.png +0 -0
  66. viztracer/web_dist/v52.0-6b9586def/assets/logo-3d.png +0 -0
  67. viztracer/web_dist/v52.0-6b9586def/assets/rec_atrace.png +0 -0
  68. viztracer/web_dist/v52.0-6b9586def/assets/rec_battery_counters.png +0 -0
  69. viztracer/web_dist/v52.0-6b9586def/assets/rec_board_voltage.png +0 -0
  70. viztracer/web_dist/v52.0-6b9586def/assets/rec_cpu_coarse.png +0 -0
  71. viztracer/web_dist/v52.0-6b9586def/assets/rec_cpu_fine.png +0 -0
  72. viztracer/web_dist/v52.0-6b9586def/assets/rec_cpu_freq.png +0 -0
  73. viztracer/web_dist/v52.0-6b9586def/assets/rec_cpu_voltage.png +0 -0
  74. viztracer/web_dist/v52.0-6b9586def/assets/rec_frame_timeline.png +0 -0
  75. viztracer/web_dist/v52.0-6b9586def/assets/rec_ftrace.png +0 -0
  76. viztracer/web_dist/v52.0-6b9586def/assets/rec_gpu_mem_total.png +0 -0
  77. viztracer/web_dist/v52.0-6b9586def/assets/rec_java_heap_dump.png +0 -0
  78. viztracer/web_dist/v52.0-6b9586def/assets/rec_lmk.png +0 -0
  79. viztracer/web_dist/v52.0-6b9586def/assets/rec_logcat.png +0 -0
  80. viztracer/web_dist/v52.0-6b9586def/assets/rec_long_trace.png +0 -0
  81. viztracer/web_dist/v52.0-6b9586def/assets/rec_mem_hifreq.png +0 -0
  82. viztracer/web_dist/v52.0-6b9586def/assets/rec_meminfo.png +0 -0
  83. viztracer/web_dist/v52.0-6b9586def/assets/rec_native_heap_profiler.png +0 -0
  84. viztracer/web_dist/v52.0-6b9586def/assets/rec_one_shot.png +0 -0
  85. viztracer/web_dist/v52.0-6b9586def/assets/rec_profiling.png +0 -0
  86. viztracer/web_dist/v52.0-6b9586def/assets/rec_ps_stats.png +0 -0
  87. viztracer/web_dist/v52.0-6b9586def/assets/rec_ring_buf.png +0 -0
  88. viztracer/web_dist/v52.0-6b9586def/assets/rec_syscalls.png +0 -0
  89. viztracer/web_dist/v52.0-6b9586def/assets/rec_vmstat.png +0 -0
  90. viztracer/web_dist/v52.0-6b9586def/assets/scheduling_latency.png +0 -0
  91. viztracer/web_dist/v52.0-6b9586def/assets/vscode-icon.png +0 -0
  92. viztracer/web_dist/v52.0-6b9586def/engine_bundle.js +3 -0
  93. viztracer/web_dist/v52.0-6b9586def/frontend_bundle.js +5495 -0
  94. viztracer/web_dist/v52.0-6b9586def/index.html +127 -0
  95. viztracer/web_dist/v52.0-6b9586def/manifest.json +52 -0
  96. viztracer/web_dist/v52.0-6b9586def/perfetto.css +5737 -0
  97. viztracer/web_dist/v52.0-6b9586def/stdlib_docs.json +1 -0
  98. viztracer/web_dist/v52.0-6b9586def/trace_config_utils.wasm +0 -0
  99. viztracer/web_dist/v52.0-6b9586def/trace_processor.wasm +0 -0
  100. viztracer/web_dist/v52.0-6b9586def/trace_processor_memory64.wasm +0 -0
  101. viztracer/web_dist/v52.0-6b9586def/traceconv.wasm +0 -0
  102. viztracer/web_dist/v52.0-6b9586def/traceconv_bundle.js +2 -0
  103. viztracer-1.1.0.dist-info/METADATA +316 -0
  104. viztracer-1.1.0.dist-info/RECORD +109 -0
  105. viztracer-1.1.0.dist-info/WHEEL +5 -0
  106. viztracer-1.1.0.dist-info/entry_points.txt +3 -0
  107. viztracer-1.1.0.dist-info/licenses/LICENSE +222 -0
  108. viztracer-1.1.0.dist-info/licenses/NOTICE.txt +27 -0
  109. viztracer-1.1.0.dist-info/top_level.txt +1 -0
viztracer/cellmagic.py ADDED
@@ -0,0 +1,70 @@
1
+ # Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
2
+ # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
3
+
4
+
5
+ def load_ipython_extension(ipython) -> None:
6
+ """
7
+ Use `%load_ext viztracer`
8
+ Lazy execute the following code because importing IPython is slow
9
+ """
10
+ from IPython.core.magic import (Magics, cell_magic, # type: ignore
11
+ magics_class, needs_local_scope)
12
+ from IPython.core.magic_arguments import (argument, magic_arguments, parse_argstring) # type: ignore
13
+
14
+ @magics_class
15
+ class VizTracerMagics(Magics):
16
+ @magic_arguments()
17
+ @argument("--port", "-p", default=9001, type=int,
18
+ help="specify the port vizviewer will use")
19
+ @argument("--output_file", default="./viztracer_report.json",
20
+ help="output file path. End with .json or .html or .gz")
21
+ @argument("--max_stack_depth", type=int, default=-1,
22
+ help="maximum stack depth you want to trace.")
23
+ @argument("--ignore_c_function", action="store_true", default=False,
24
+ help="ignore all c functions including most builtin functions and libraries")
25
+ @argument("--ignore_frozen", action="store_true", default=False,
26
+ help="ignore all functions that are frozen(like import)")
27
+ @argument("--log_func_args", action="store_true", default=False,
28
+ help="log all function arguments, this will introduce large overhead")
29
+ @argument("--log_print", action="store_true", default=False,
30
+ help="replace all print() function to adding an event to the result")
31
+ @argument("--log_sparse", action="store_true", default=False,
32
+ help="log only selected functions with @log_sparse")
33
+ @needs_local_scope
34
+ @cell_magic
35
+ def viztracer(self, line, cell, local_ns) -> None:
36
+ from IPython.display import display # type: ignore
37
+ from ipywidgets import Button # type: ignore
38
+
39
+ from .viewer import ServerThread
40
+ from .viztracer import VizTracer
41
+ options = parse_argstring(self.viztracer, line)
42
+ assert self.shell is not None
43
+ code = self.shell.transform_cell(cell)
44
+ file_path = options.output_file
45
+ tracer_kwargs = {
46
+ "output_file": file_path,
47
+ "verbose": 0,
48
+ "max_stack_depth": options.max_stack_depth,
49
+ "ignore_c_function": options.ignore_c_function,
50
+ "ignore_frozen": options.ignore_frozen,
51
+ "log_func_args": options.log_func_args,
52
+ "log_print": options.log_print,
53
+ "log_sparse": options.log_sparse
54
+ }
55
+ with VizTracer(**tracer_kwargs):
56
+ exec(code, local_ns, local_ns)
57
+
58
+ def view(): # pragma: no cover
59
+ server = ServerThread(file_path, port=options.port, once=True)
60
+ server.start()
61
+ server.ready.wait()
62
+ import webbrowser
63
+ webbrowser.open_new_tab(f'http://127.0.0.1:{server.port}')
64
+
65
+ button = Button(description="VizTracer Report")
66
+ button.on_click(lambda b: view())
67
+
68
+ display(button)
69
+
70
+ ipython.register_magics(VizTracerMagics)
@@ -0,0 +1,353 @@
1
+ # Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
2
+ # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
3
+
4
+ import ast
5
+ import copy
6
+ import re
7
+ from functools import reduce
8
+ from typing import Any, Callable
9
+
10
+ from .util import color_print
11
+
12
+
13
+ class AstTransformer(ast.NodeTransformer):
14
+ def __init__(self, inst_type: str, inst_args: dict[str, dict]) -> None:
15
+ super().__init__()
16
+ self.inst_type: str = inst_type
17
+ self.inst_args: dict[str, dict] = inst_args
18
+ self.curr_lineno: int = 0
19
+ self.log_func_exec_enable: bool = False
20
+
21
+ def visit_Assign(self, node: ast.Assign) -> ast.stmt | list[ast.stmt]:
22
+ return self._visit_generic_assign(node)
23
+
24
+ def visit_AugAssign(self, node: ast.AugAssign) -> ast.stmt | list[ast.stmt]:
25
+ return self._visit_generic_assign(node)
26
+
27
+ def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.stmt | list[ast.stmt]:
28
+ return self._visit_generic_assign(node)
29
+
30
+ def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
31
+ if self.inst_type == "log_func_exec":
32
+ for funcname in self.inst_args["funcnames"]:
33
+ if re.fullmatch(funcname, node.name):
34
+ self.log_func_exec_enable = True
35
+ elif self.inst_type == "log_func_entry":
36
+ for funcname in self.inst_args["funcnames"]:
37
+ if re.fullmatch(funcname, node.name):
38
+ node.body.insert(0, self.get_instrument_node("Function Entry", node.name))
39
+ elif self.inst_type in ("log_var", "log_number"):
40
+ instrumented_nodes: list[ast.stmt] = []
41
+ args = node.args
42
+ func_args_name = [a.arg for a in args.posonlyargs + args.args + args.kwonlyargs]
43
+ if "vararg" in args._fields and args.vararg:
44
+ func_args_name.append(args.vararg.arg)
45
+ if "kwarg" in args._fields and args.kwarg:
46
+ func_args_name.append(args.kwarg.arg)
47
+ for name in func_args_name:
48
+ for pattern in self.inst_args["varnames"]:
49
+ if re.fullmatch(pattern, name):
50
+ instrumented_nodes.append(self.get_instrument_node("Variable Assign", name))
51
+ break
52
+
53
+ self.generic_visit(node)
54
+
55
+ if self.inst_type == "log_func_exec":
56
+ self.log_func_exec_enable = False
57
+ elif self.inst_type in ("log_var", "log_number") and instrumented_nodes:
58
+ node.body = instrumented_nodes + node.body
59
+ return node
60
+
61
+ def visit_For(self, node: ast.For) -> ast.For:
62
+ if self.inst_type in ("log_var", "log_number"):
63
+ instrumented_nodes = self.get_assign_log_nodes(node.target)
64
+
65
+ self.generic_visit(node)
66
+
67
+ if self.inst_type in ("log_var", "log_number"):
68
+ if instrumented_nodes:
69
+ node.body = instrumented_nodes + node.body
70
+ return node
71
+
72
+ def visit_Raise(self, node: ast.Raise) -> ast.AST | list[ast.AST]:
73
+ if self.inst_type == "log_exception":
74
+ instrument_node = self.get_instrument_node_by_node("Exception", node.exc)
75
+ return [instrument_node, node]
76
+ return node
77
+
78
+ def _visit_generic_assign(self, node: ast.Assign | ast.AugAssign | ast.AnnAssign) -> list[ast.stmt]:
79
+ self.generic_visit(node)
80
+ ret: list[ast.stmt] = [node]
81
+ self.curr_lineno = node.lineno
82
+ if self.inst_type in ("log_var", "log_number", "log_attr", "log_func_exec"):
83
+ if isinstance(node, (ast.AugAssign, ast.AnnAssign)):
84
+ instrumented_nodes = self.get_assign_log_nodes(node.target)
85
+ if instrumented_nodes:
86
+ ret.extend(instrumented_nodes)
87
+ elif isinstance(node, ast.Assign):
88
+ for target in node.targets:
89
+ instrumented_nodes = self.get_assign_log_nodes(target)
90
+ if instrumented_nodes:
91
+ ret.extend(instrumented_nodes)
92
+ return ret
93
+
94
+ def get_assign_targets(self, node: ast.expr) -> list[str]:
95
+ """
96
+ :param ast.Node node:
97
+ """
98
+ if isinstance(node, ast.Name):
99
+ return [node.id]
100
+ elif isinstance(node, (ast.Attribute, ast.Subscript, ast.Starred)):
101
+ return self.get_assign_targets(node.value)
102
+ elif isinstance(node, ast.Tuple) or isinstance(node, ast.List):
103
+ return reduce(lambda a, b: a + b, [self.get_assign_targets(elt) for elt in node.elts])
104
+ color_print("WARNING", "Unexpected node type {} for ast.Assign. \
105
+ Please report to the author github.com/gaogaotiantian/viztracer".format(type(node)))
106
+ return []
107
+
108
+ def get_assign_targets_with_attr(self, node: ast.AST) -> list[ast.Attribute]:
109
+ """
110
+ :param ast.Node node:
111
+ """
112
+ if isinstance(node, ast.Attribute):
113
+ return [node]
114
+ elif isinstance(node, (ast.Name, ast.Subscript, ast.Starred)):
115
+ return []
116
+ elif isinstance(node, (ast.Tuple, ast.List)):
117
+ return reduce(lambda a, b: a + b, [self.get_assign_targets_with_attr(elt) for elt in node.elts])
118
+ color_print("WARNING", "Unexpected node type {} for ast.Assign. \
119
+ Please report to the author github.com/gaogaotiantian/viztracer".format(type(node)))
120
+ return []
121
+
122
+ def get_assign_log_nodes(self, target: ast.expr) -> list[ast.stmt]:
123
+ """
124
+ given a target of any type of Assign, return the instrumented node
125
+ that log this variable
126
+ if this target is not supposed to be logged, return []
127
+ """
128
+ ret: list[ast.stmt] = []
129
+ if self.inst_type in ("log_var", "log_number"):
130
+ target_ids = self.get_assign_targets(target)
131
+ for target_id in target_ids:
132
+ for varname in self.inst_args["varnames"]:
133
+ if re.fullmatch(varname, target_id):
134
+ ret.append(self.get_instrument_node("Variable Assign", target_id))
135
+ break
136
+ elif self.inst_type == "log_attr":
137
+ target_nodes = self.get_assign_targets_with_attr(target)
138
+ for target_node in target_nodes:
139
+ for varname in self.inst_args["varnames"]:
140
+ if re.fullmatch(varname, target_node.attr):
141
+ ret.append(self.get_instrument_node_by_node("Attribute Assign", target_node))
142
+ break
143
+ elif self.inst_type == "log_func_exec":
144
+ if self.log_func_exec_enable:
145
+ target_ids = self.get_assign_targets(target)
146
+ for target_id in target_ids:
147
+ ret.append(self.get_instrument_node("Variable Assign", target_id))
148
+
149
+ return ret
150
+
151
+ def get_instrument_node(self, trigger: str, name: str) -> ast.Expr:
152
+ if self.inst_type in ("log_var", "log_number"):
153
+ if self.inst_type == "log_var":
154
+ event = "instant"
155
+ elif self.inst_type == "log_number":
156
+ event = "counter"
157
+ return self.get_add_variable_node(
158
+ name=f"{trigger} - {name}",
159
+ var_node=ast.Name(id=name, ctx=ast.Load()),
160
+ event=event,
161
+ )
162
+ elif self.inst_type == "log_func_exec":
163
+ return self.get_add_func_exec_node(
164
+ name=f"{name}",
165
+ val=ast.Name(id=name, ctx=ast.Load()),
166
+ lineno=self.curr_lineno,
167
+ )
168
+ elif self.inst_type == "log_func_entry":
169
+ return self.get_add_variable_node(
170
+ name=f"{trigger} - {name}",
171
+ var_node=ast.Constant(value=f"{name} is called"),
172
+ event="instant",
173
+ )
174
+ else:
175
+ raise ValueError(f"{name} is not supported")
176
+
177
+ def get_instrument_node_by_node(self, trigger: str, node: ast.expr | None) -> ast.Expr:
178
+ var_node: ast.expr
179
+ if node is None:
180
+ name = f"{trigger}"
181
+ var_node = ast.Constant(value=None)
182
+ else:
183
+ name = f"{trigger} - {self.get_string_of_expr(node)}"
184
+ var_node = self.copy_node_with_load(node)
185
+ return self.get_add_variable_node(
186
+ name=name,
187
+ var_node=var_node,
188
+ event="instant",
189
+ )
190
+
191
+ def get_add_variable_node(self, name: str, var_node: ast.expr, event: str) -> ast.Expr:
192
+ node_instrument = ast.Expr(
193
+ value=ast.Call(
194
+ func=ast.Attribute(
195
+ value=ast.Name(id="__viz_tracer__", ctx=ast.Load()),
196
+ attr="add_variable",
197
+ ctx=ast.Load(),
198
+ ),
199
+ args=[
200
+ ast.Constant(value=name),
201
+ var_node,
202
+ ast.Constant(value=event),
203
+ ],
204
+ keywords=[],
205
+ ),
206
+ )
207
+ return node_instrument
208
+
209
+ def get_add_func_exec_node(self, name: str, val: ast.AST, lineno: int) -> ast.Expr:
210
+ node_instrument = ast.Expr(
211
+ value=ast.Call(
212
+ func=ast.Attribute(
213
+ value=ast.Name(id="__viz_tracer__", ctx=ast.Load()),
214
+ attr="add_func_exec",
215
+ ctx=ast.Load(),
216
+ ),
217
+ args=[
218
+ ast.Constant(value=name),
219
+ ast.Name(id=name, ctx=ast.Load()),
220
+ ast.Constant(value=lineno),
221
+ ],
222
+ keywords=[],
223
+ ),
224
+ )
225
+ return node_instrument
226
+
227
+ def copy_node_with_load(self, node: ast.expr) -> ast.expr:
228
+ """
229
+ copy the whole node tree but change all Store to Load
230
+ """
231
+ new_node = copy.deepcopy(node)
232
+ for n in ast.walk(new_node):
233
+ # Fix Store to Load
234
+ if "ctx" in n._fields and isinstance(n.ctx, ast.Store): # type: ignore
235
+ n.ctx = ast.Load() # type: ignore
236
+ return new_node
237
+
238
+ def get_string_of_expr(self, node: ast.expr | ast.slice) -> str:
239
+ """
240
+ Try to do "unparse" of the node
241
+ """
242
+ if isinstance(node, ast.Name):
243
+ return node.id
244
+ elif isinstance(node, ast.Constant):
245
+ return repr(node.value)
246
+ elif isinstance(node, ast.Attribute):
247
+ return f"{self.get_string_of_expr(node.value)}.{node.attr}"
248
+ elif isinstance(node, ast.Subscript):
249
+ return f"{self.get_string_of_expr(node.value)}[{self.get_string_of_expr(node.slice)}]"
250
+ elif isinstance(node, ast.Call):
251
+ return f"{self.get_string_of_expr(node.func)}()"
252
+ elif isinstance(node, ast.Starred):
253
+ return f"*{self.get_string_of_expr(node.value)}"
254
+ elif isinstance(node, ast.Tuple):
255
+ return f"({','.join([self.get_string_of_expr(elt) for elt in node.elts])})"
256
+ elif isinstance(node, ast.List):
257
+ return f"[{','.join([self.get_string_of_expr(elt) for elt in node.elts])}]"
258
+ elif isinstance(node, ast.Slice):
259
+ lower = self.get_string_of_expr(node.lower) if "lower" in node._fields and node.lower else ""
260
+ upper = self.get_string_of_expr(node.upper) if "upper" in node._fields and node.upper else ""
261
+ step = self.get_string_of_expr(node.step) if "step" in node._fields and node.step else ""
262
+ if step:
263
+ return f"{lower}:{upper}:{step}"
264
+ elif upper:
265
+ return f"{lower}:{upper}"
266
+ else:
267
+ return f"{lower}:"
268
+ color_print("WARNING", "Unexpected node type {} for ast.Assign. \
269
+ Please report to the author github.com/gaogaotiantian/viztracer".format(type(node)))
270
+ return ""
271
+
272
+
273
+ class SourceProcessor:
274
+ """
275
+ Pre-process comments like #!viztracer: log_instant("event")
276
+ """
277
+
278
+ def process(self, source: Any):
279
+ if isinstance(source, bytes):
280
+ source = source.decode("utf-8")
281
+ elif not isinstance(source, str):
282
+ return source
283
+
284
+ new_lines = []
285
+
286
+ for line in source.splitlines():
287
+ for pattern, transform in self.re_patterns:
288
+ m = pattern.match(line)
289
+ if m:
290
+ new_lines.append(transform(self, m))
291
+ break
292
+ else:
293
+ new_lines.append(line)
294
+
295
+ return "\n".join(new_lines)
296
+
297
+ def line_transform(self, re_match: re.Match) -> str:
298
+ return f"{re_match.group(1)}__viz_tracer__.{re_match.group(2)}"
299
+
300
+ def line_transform_condition(self, re_match: re.Match) -> str:
301
+ return f"{re_match.group(1)}if {re_match.group(3)}: __viz_tracer__.{re_match.group(2)};"
302
+
303
+ def inline_transform(self, re_match: re.Match) -> str:
304
+ stmt = re_match.group(1)
305
+ if "=" in stmt:
306
+ val_assigned = stmt[:stmt.index("=")].strip()
307
+ return f"{stmt}; __viz_tracer__.log_var('{val_assigned}', ({val_assigned}))"
308
+ return f"{stmt}; __viz_tracer__.log_instant('{stmt.strip()}')"
309
+
310
+ def inline_transform_condition(self, re_match: re.Match) -> str:
311
+ stmt = re_match.group(1)
312
+ if "=" in stmt:
313
+ val_assigned = stmt[:stmt.index("=")].strip()
314
+ return f"{stmt}; __viz_tracer__.log_var('{val_assigned}', ({val_assigned}), cond={re_match.group(2)})"
315
+ return f"{stmt}; __viz_tracer__.log_instant('{stmt.strip()}', cond={re_match.group(2)});"
316
+
317
+ re_patterns = [
318
+ # !viztracer: log_var("var", var)
319
+ (re.compile(r"(\s*)#\s*!viztracer:\s*(log_.*?\(.*\))\s*$"), line_transform),
320
+ # a = 3 # !viztracer: log
321
+ (re.compile(r"(.*\S.*)#\s*!viztracer:\s*log\s*$"), inline_transform),
322
+ # !viztracer: log_var("var", var) if var > 3
323
+ (re.compile(r"(\s*)#\s*!viztracer:\s*(log_.*?\(.*\))\s*if\s+(.*?)\s*$"), line_transform_condition),
324
+ # a = 3 # !viztracer: log if a != 3
325
+ (re.compile(r"(.*\S.*)#\s*!viztracer:\s*log\s*if\s+(.*?)\s*$"), inline_transform_condition),
326
+ ]
327
+
328
+
329
+ class CodeMonkey:
330
+ def __init__(self, file_name: str) -> None:
331
+ self.file_name: str = file_name
332
+ self._compile: Callable = compile
333
+ self.source_processor: SourceProcessor | None = None
334
+ self.ast_transformers: list[AstTransformer] = []
335
+
336
+ def add_instrument(self, inst_type: str, inst_args: dict[str, dict]) -> None:
337
+ self.ast_transformers.append(AstTransformer(inst_type, inst_args))
338
+
339
+ def add_source_processor(self) -> None:
340
+ self.source_processor = SourceProcessor()
341
+
342
+ def compile(self, source, filename, mode, flags=0, dont_inherit=False, optimize=-1, *, _feature_version=-1):
343
+ if self.source_processor is not None:
344
+ source = self.source_processor.process(source)
345
+ if self.ast_transformers:
346
+ tree = self._compile(source, filename, mode, flags | ast.PyCF_ONLY_AST,
347
+ dont_inherit, optimize, _feature_version=_feature_version)
348
+ for trans in self.ast_transformers:
349
+ trans.visit(tree)
350
+ ast.fix_missing_locations(tree)
351
+ return self._compile(tree, filename, mode, flags, dont_inherit, optimize, _feature_version=_feature_version)
352
+
353
+ return self._compile(source, filename, mode, flags, dont_inherit, optimize, _feature_version=_feature_version)
viztracer/decorator.py ADDED
@@ -0,0 +1,164 @@
1
+ # Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
2
+ # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
3
+
4
+ import functools
5
+ import multiprocessing
6
+ import os
7
+ import time
8
+ from typing import Any, Callable, TypeVar, overload
9
+
10
+ from .viztracer import VizTracer, get_tracer
11
+
12
+
13
+ R = TypeVar("R")
14
+
15
+
16
+ @overload
17
+ def ignore_function(method: None,
18
+ tracer: VizTracer | None = None) -> Callable[[Callable[..., R]], Callable[..., R]]:
19
+ pass # pragma: no cover
20
+
21
+
22
+ @overload
23
+ def ignore_function(method: Callable[..., R],
24
+ tracer: VizTracer | None = None) -> Callable[..., R]:
25
+ pass # pragma: no cover
26
+
27
+
28
+ def ignore_function(method: Callable[..., R] | None = None,
29
+ tracer: VizTracer | None = None) -> Callable[..., R] | Callable[[Callable[..., R]], Callable[..., R]]:
30
+
31
+ def inner(func: Callable[..., R]) -> Callable[..., R]:
32
+
33
+ @functools.wraps(func)
34
+ def ignore_wrapper(*args, **kwargs) -> Any:
35
+ # We need this to keep trace a local variable
36
+ t = tracer
37
+ if not t:
38
+ t = get_tracer()
39
+ if not t:
40
+ raise NameError("ignore_function only works with global tracer")
41
+ t.pause()
42
+ ret = func(*args, **kwargs)
43
+ t.resume()
44
+ return ret
45
+
46
+ return ignore_wrapper
47
+
48
+ if method:
49
+ return inner(method)
50
+ return inner
51
+
52
+
53
+ @overload
54
+ def trace_and_save(method: None,
55
+ output_dir: str = "./",
56
+ **viztracer_kwargs) -> Callable[[Callable[..., R]], Callable[..., R]]:
57
+ pass # pragma: no cover
58
+
59
+
60
+ @overload
61
+ def trace_and_save(method: Callable[..., R],
62
+ output_dir: str = "./",
63
+ **viztracer_kwargs) -> Callable[..., R]:
64
+ pass # pragma: no cover
65
+
66
+
67
+ def trace_and_save(method: Callable[..., R] | None = None,
68
+ output_dir: str = "./",
69
+ **viztracer_kwargs) -> Callable[..., R] | Callable[[Callable[..., R]], Callable[..., R]]:
70
+
71
+ def inner(func: Callable[..., R]) -> Callable[..., R]:
72
+
73
+ @functools.wraps(func)
74
+ def wrapper(*args, **kwargs) -> Any:
75
+ tracer = VizTracer(**viztracer_kwargs)
76
+ tracer.start()
77
+ ret = func(*args, **kwargs)
78
+ tracer.stop()
79
+ if not os.path.exists(output_dir):
80
+ os.mkdir(output_dir)
81
+ file_name = os.path.join(output_dir, f"result_{func.__name__}_{int(100000 * time.time())}.json")
82
+ if multiprocessing.get_start_method() == "fork" and not multiprocessing.current_process().daemon:
83
+ tracer.fork_save(file_name)
84
+ else:
85
+ tracer.save(file_name)
86
+ tracer.clear()
87
+ return ret
88
+
89
+ return wrapper
90
+
91
+ if method:
92
+ return inner(method)
93
+ return inner
94
+
95
+
96
+ def _log_sparse_wrapper(func: Callable, stack_depth: int = 0,
97
+ dynamic_tracer_check: bool = False) -> Callable:
98
+ if not dynamic_tracer_check:
99
+ tracer = get_tracer()
100
+ if tracer is None or not tracer.log_sparse:
101
+ return func
102
+
103
+ @functools.wraps(func)
104
+ def wrapper(*args, **kwargs) -> Any:
105
+ local_tracer = get_tracer() if dynamic_tracer_check else tracer
106
+
107
+ if local_tracer is None:
108
+ return func(*args, **kwargs)
109
+ assert isinstance(local_tracer, VizTracer)
110
+
111
+ if local_tracer.log_sparse and not local_tracer.enable:
112
+ if stack_depth > 0:
113
+ orig_max_stack_depth = local_tracer.max_stack_depth
114
+ local_tracer.max_stack_depth = stack_depth
115
+ local_tracer.start()
116
+ ret = func(*args, **kwargs)
117
+ local_tracer.stop()
118
+ local_tracer.max_stack_depth = orig_max_stack_depth
119
+ return ret
120
+ else:
121
+ start = local_tracer.getts()
122
+ ret = func(*args, **kwargs)
123
+ dur = local_tracer.getts() - start
124
+ code = func.__code__
125
+ raw_data = {
126
+ "ph": "X",
127
+ "name": f"{code.co_name} ({code.co_filename}:{code.co_firstlineno})",
128
+ "ts": start,
129
+ "dur": dur,
130
+ "cat": "FEE",
131
+ }
132
+ local_tracer.add_raw(raw_data)
133
+ return ret
134
+ elif local_tracer.enable and not local_tracer.log_sparse:
135
+ # The call is made from the module inside, so if `trace_self=False` it will be ignored.
136
+ # To avoid this behavior, we need to reset the counter `ignore_stack_depth`` and then
137
+ # recover it
138
+ return local_tracer.shield_ignore(func, *args, **kwargs)
139
+ else:
140
+ return func(*args, **kwargs)
141
+
142
+ return wrapper
143
+
144
+
145
+ @overload
146
+ def log_sparse(func: None,
147
+ stack_depth: int = 0,
148
+ dynamic_tracer_check: bool = False) -> Callable[[Callable[..., R]], Callable[..., R]]:
149
+ pass # pragma: no cover
150
+
151
+
152
+ @overload
153
+ def log_sparse(func: Callable[..., R],
154
+ stack_depth: int = 0,
155
+ dynamic_tracer_check: bool = False) -> Callable[..., R]:
156
+ pass # pragma: no cover
157
+
158
+
159
+ def log_sparse(func: Callable[..., R] | None = None,
160
+ stack_depth: int = 0,
161
+ dynamic_tracer_check: bool = False) -> Callable[..., R] | Callable[[Callable[..., R]], Callable[..., R]]:
162
+ if func is None:
163
+ return functools.partial(_log_sparse_wrapper, stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)
164
+ return _log_sparse_wrapper(func=func, stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)
@@ -0,0 +1,81 @@
1
+ # Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
2
+ # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
3
+
4
+ import functools
5
+ from typing import Any, Callable, Literal
6
+
7
+ from .viztracer import VizTracer
8
+
9
+
10
+ class _EventBase:
11
+ def __init__(self, tracer: VizTracer, name: str = "", **kwargs) -> None:
12
+ self._viztracer_tracer: VizTracer = tracer
13
+ self._viztracer_name: str = name
14
+ self._viztracer_enable: bool = False
15
+ self._viztracer_config: dict = {
16
+ "trigger_on_change": True,
17
+ "include_attributes": [],
18
+ "exclude_attributes": [],
19
+ }
20
+
21
+ for key in kwargs:
22
+ if key in self._viztracer_config:
23
+ self._viztracer_config[key] = kwargs[key]
24
+
25
+ self._viztracer_enable = True
26
+
27
+ def __setattr__(self, name: str, value: Any) -> None:
28
+ self.__dict__[name] = value
29
+ if not name.startswith("_"):
30
+ if self._viztracer_enable and self._viztracer_config["trigger_on_change"]:
31
+ if self._viztracer_config["include_attributes"]:
32
+ if name in self._viztracer_config["include_attributes"]:
33
+ self._viztracer_log()
34
+ elif self._viztracer_config["exclude_attributes"]:
35
+ if name not in self._viztracer_config["exclude_attributes"]:
36
+ self._viztracer_log()
37
+ else:
38
+ self._viztracer_log()
39
+
40
+ def _viztracer_get_attr_list(self) -> list[str]:
41
+ if self._viztracer_config["include_attributes"]:
42
+ return self._viztracer_config["include_attributes"]
43
+ else:
44
+ return [attr for attr in self.__dir__()
45
+ if not attr.startswith("_") and attr not in self._viztracer_config["exclude_attributes"]]
46
+
47
+ def _viztracer_set_config(self, key: str, value: Any) -> None:
48
+ if key not in self._viztracer_config:
49
+ raise ValueError(f"No config named {key}")
50
+ self._viztracer_config[key] = value
51
+
52
+ def config(self, key: str, value: Any) -> None:
53
+ self._viztracer_set_config(key, value)
54
+
55
+ def _viztracer_log(self) -> None:
56
+ raise NotImplementedError("You should not use _EventBase class directly")
57
+
58
+ def log(self) -> None:
59
+ self._viztracer_log()
60
+
61
+ @staticmethod
62
+ def triggerlog(method: Callable | None = None,
63
+ when: Literal["after", "before", "both"] = "after") -> Callable:
64
+ if when not in ["after", "before", "both"]:
65
+ raise ValueError(f"when has to be one of 'after', 'before' or 'both', not {when}")
66
+
67
+ def inner(func: Callable) -> Callable:
68
+
69
+ @functools.wraps(func)
70
+ def wrapper(self, *args, **kwargs) -> Any:
71
+ if when in ("before", "both"):
72
+ self._viztracer_log()
73
+ ret = func(self, *args, **kwargs)
74
+ if when in ("after", "both"):
75
+ self._viztracer_log()
76
+ return ret
77
+ return wrapper
78
+
79
+ if method:
80
+ return inner(method)
81
+ return inner