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