1# Copyright 2025 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""func tracer module"""
16import ast
17import sys
18import inspect
19import textwrap
20import linecache
21
22
23class _FuncTracer:
24 """Func tracer class"""
25
26 def __init__(self):
27 self.ind = 0
28 self.trace_str = ""
29 self.code_trees = {}
30 self.max_sub_length = 30
31
32 def add2trace(self, s):
33 """log"""
34 self.trace_str += "\t" * self.ind + s + "\n"
35
36 def is_constant_expr(self, s):
37 """whether it's a leave node"""
38 node = ast.parse(s, mode="eval").body
39 return isinstance(node, (ast.Constant, ast.Attribute, ast.Name))
40
41 def scrap_term_symbols(self, s):
42 """get all leave nodes from source code"""
43 res = []
44 tree = ast.parse(s)
45 for node in ast.walk(tree):
46 if isinstance(node, ast.Name):
47 if not any(r.startswith(node.id) for r in res):
48 res += [node.id]
49 elif isinstance(node, ast.Constant):
50 res += [str(node.value)]
51 elif isinstance(node, ast.Attribute):
52 # Reconstruct full attribute chain like obj.attr1.attr2
53 value = node
54 attr_chain = []
55 while isinstance(value, ast.Attribute):
56 attr_chain.insert(0, value.attr)
57 value = value.value
58 if isinstance(value, ast.Name):
59 attr_chain.insert(0, value.id)
60 res += [".".join(attr_chain)]
61 res = sorted(res, key=len, reverse=True)
62 return res
63
64 def substitute(self, s, local):
65 """substitute leave nodes by their numerical values"""
66 nodes = self.scrap_term_symbols(s)
67 for n in nodes:
68 attr_chain = n.split(".")
69 if attr_chain[0] in local:
70 if len(attr_chain) == 1:
71 target = str(local[n])
72 if len(target) < self.max_sub_length:
73 s = s.replace(n, target)
74 else:
75 s = s.replace(n, type(local[n]).__name__)
76 else:
77 attr = local[attr_chain[0]]
78 for a in attr_chain[1:]:
79 attr = getattr(attr, a)
80 target = str(attr)
81 if len(target) < self.max_sub_length:
82 s = s.replace(n, target)
83 else:
84 s = s.replace(n, type(attr).__name__)
85 return s
86
87 def fetch_node_from_lineno(self, lineno, co):
88 """Get AST node from a line num in source code"""
89 for node in ast.walk(self.code_trees[co]):
90 if hasattr(node, "lineno") and node.lineno == lineno:
91 return node
92 return None
93
94 def line_tracer(self, frame, event, _):
95 """Trace executed code line"""
96 co = frame.f_code
97 func_name = co.co_name
98
99 # Handling function call
100 if event == "call":
101 self.extract_ast(co)
102 _, _, _, values = inspect.getargvalues(frame)
103 args_str = ",".join(
104 f"{k}={v if isinstance(v, (int, float)) else type(v).__name__}"
105 for k, v in values.items()
106 )
107 self.add2trace(f"::{func_name}({args_str})")
108 self.ind += 1
109
110 # Handling instruction
111 elif event == "line":
112 line = linecache.getline(co.co_filename, frame.f_lineno).strip()
113 ignored_statements = ("if ", "elif ", "else ", "for ", "while ")
114 op_equals = ["*=", "+=", "-=", "/=", "%=", "="]
115 # Only capture assignments and returns
116 if not line.startswith(ignored_statements):
117 base_str = f"{func_name}->"
118 if "=" in line:
119 try:
120 ast.literal_eval(line)
121 except (ValueError, SyntaxError):
122 node = self.fetch_node_from_lineno(frame.f_lineno, co)
123 if not node:
124 return self.line_tracer
125 line = ast.unparse(node)
126 sign = next(op for op in op_equals if op in line)
127 left, right = line.split(sign, 1)
128 left, right = left.strip(), right.strip()
129 if sign != "=":
130 right = f"{left} {sign[0]} {right}"
131 self.add2trace(f"{base_str} {left} = {right}")
132 if not self.is_constant_expr(right):
133 val = eval( # pylint: disable=eval-used
134 right, frame.f_globals, frame.f_locals
135 )
136 sub_right = self.substitute(right, frame.f_locals)
137 base_str = f"{base_str} {left}"
138 if not self.is_constant_expr(sub_right):
139 self.add2trace(
140 f"{len(base_str)*' '} = {sub_right}"
141 )
142 self.add2trace(f"{len(base_str)*' '} = {val}")
143 else:
144 self.add2trace(f"({right}) = {val}")
145 # Handle return statements
146 elif line.startswith("return"):
147 lreturn = len("return")
148 expr = line[lreturn:].strip()
149 try:
150 ast.literal_eval(expr)
151 except (ValueError, SyntaxError):
152 node = self.fetch_node_from_lineno(frame.f_lineno, co)
153 if not node:
154 return self.line_tracer
155 expr = ast.unparse(node)[lreturn:].strip()
156 val = eval( # pylint: disable=eval-used
157 expr, frame.f_globals, frame.f_locals
158 )
159 sub_expr = self.substitute(expr, frame.f_locals)
160 self.add2trace(f"{base_str} return {expr}")
161 if not self.is_constant_expr(sub_expr):
162 self.add2trace(f"{len(base_str)*' '} = {sub_expr}")
163 self.add2trace(f"{len(base_str)*' '} = {val}")
164 else:
165 self.add2trace(f"{len(base_str)*' '} = {val}")
166
167 elif event == "return":
168 self.ind -= 1
169
170 return self.line_tracer
171
172 def extract_ast(self, fun):
173 """Get AST from function source code"""
174 if fun not in self.code_trees:
175 self.code_trees[fun] = ast.parse(
176 textwrap.dedent(inspect.getsource(fun))
177 )
178 _, start_line = inspect.getsourcelines(fun)
179 # Extract file's real line numbers for AST
180 for node in ast.walk(self.code_trees[fun]):
181 if hasattr(node, "lineno"):
182 node.lineno = start_line + node.lineno - 1
183 node.end_lineno = start_line + node.end_lineno - 1
184
185 def wrap(self, fun):
186 """Wrapper"""
187
188 def tracked_fun(*args, **kwargs):
189 """Wrapper"""
190 sys.settrace(self.line_tracer)
191 try:
192 print("Tracing...")
193 res = fun(*args, **kwargs)
194 print(self.trace_str)
195 self.trace_str = ""
196 return res
197 finally:
198 sys.settrace(None)
199 return res
200
201 return tracked_fun