Skip to content

fix(ptodsl): preserve original source line numbers after AST rewrite - #1053

Open
liggest wants to merge 2 commits into
hw-native-sys:mainfrom
liggest:fix_ast_rewrite_tb
Open

fix(ptodsl): preserve original source line numbers after AST rewrite#1053
liggest wants to merge 2 commits into
hw-native-sys:mainfrom
liggest:fix_ast_rewrite_tb

Conversation

@liggest

@liggest liggest commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Fixes #1052
PTODSL @pto.jitast_rewrite 原本使用 inspect.getsource() 获取函数源码,源码被重新解析后行号从 1 开始,和原文件名一起用于编译,导致调用栈中文件名正确,行号却可能错误,少了函数在源文件中的偏移

这个 PR 改用 inspect.getsourcelines() 同时获取函数源码和函数定义的起始行,并在编译前为被解析的函数恢复行号偏移,使得调用栈中能如预期般展示出报错位置对应的正确行号和代码

测试

新增测试:

  • 触发一次编译异常,确认调用栈显示出报错函数中被标记的行
  • 确认被重新解析的函数的首行行号与原函数一致
python ptodsl/tests/test_jit_diagnostics.py
# ptodsl_jit_diagnostics: PASS
python ptodsl/tests/test_ast_rewrite_example_ir.py
# ptodsl_ast_rewrite_example_ir: PASS
python ptodsl/tests/test_jit_compile.py
# ptodsl_jit_compile: PASS

@liggest
liggest requested a review from Zhendong404 August 4, 2026 06:18

@Zhendong404 Zhendong404 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

整体方向正确:改用 inspect.getsourcelines + ast.increment_lineno 恢复行号偏移,对带装饰器的函数也是自洽的(getsourcelines 回溯到首个 @ 行,而 parse 出的 def 相对行号已包含装饰器行数),increment_lineno 同时平移 lineno/end_lineno,合成节点经 fix_missing_locations 继承父行号后也被整体平移,没有遗漏。两个新测试能证明主场景(traceback 行号正确 + co_firstlineno 一致)。可以合。

几个建议合并前或跟进处理:

  1. 改写期抛出的 PTODSLAstRewriteError 仍无位置信息(见 _ast_rewrite.py:54 的行内评论)——这是本 PR 最自然的补全,ast_rewrite 用户最常见的报错恰恰发生在改写阶段。
  2. 行为变化未声明:重写后函数的 co_firstlineno 从“片段内相对行”变为“源文件绝对行”,若下游有按旧行号做缓存 key / 快照断言 / 日志匹配的逻辑会悄悄失效,建议在 changelog 或 PR 描述里点一句。
  3. 测试缺口:只覆盖了模块级函数,真正考验 dedent + increment_lineno 组合的是嵌套/缩进的 jit 函数和多行装饰器场景,建议各补一个用例。

function_def.body = rewriter.rewrite_block(function_def.body, live_after=set())
tree = ast.Module(body=[function_def], type_ignores=[])
ast.fix_missing_locations(tree)
ast.increment_lineno(tree, source_start_line - 1)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

行号恢复本身没问题,但注意改写阶段(rewrite_block / _rewrite_if / _rewrite_for 等处)抛出的 PTODSLAstRewriteError 只带 message,没有 lineno/offset/filename。它是 SyntaxError 子类,这里树里已经有正确行号了,建议 raise 时把出错节点的 lineno 附上(可在 increment_lineno 之后再做改写,或 raise 处加 source_start_line - 1 偏移),这样“break/continue 不支持”“for-else 不支持”这类错误也能告诉用户具体是哪一行。否则这个 PR 只修了运行时 traceback,改写期诊断依然是盲区。

另外一个小遗留:textwrap.dedent 之后节点的 col_offset 是相对 dedent 后源码的。以前行号错时没人注意;现在行号对了、traceback 会显示真实源码行,一旦带 caret(SyntaxError 的 ^^^^ 指示),嵌套函数的 caret 会偏左一个缩进宽度。可以后续处理,但建议在 PR 描述里承认这个限制。

@liggest liggest Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PTODSLAstRewriteError 附上 lineno、带上有正确缩进的函数 caret
这两个可以看作独立的新特性吧,感觉有点超出这个修复的范畴了,适合后续通过新的 issue/PR 来跟踪

Comment thread ptodsl/ptodsl/_ast_rewrite.py Outdated
try:
source = inspect.getsource(fn)
source_lines, source_start_line = inspect.getsourcelines(fn)
source = "".join(source_lines)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

吹毛求疵:"".join(source_lines)inspect.getsource(fn) 完全等价,可以保留 getsource + 单独用 inspect.findsource(fn)[1] 取起始行,语义更清楚。不改也行。

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

已修改

Comment thread ptodsl/tests/test_jit_diagnostics.py Outdated
traceback_lines.append(current.tb_lineno)
current = current.tb_next
expect(
traceback_lines == [expected_line],

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

traceback_lines == [expected_line] 严格假设被测函数在栈中恰好出现一帧。将来若 rewrite 或 jit 包装引入间接调用层(多出一帧或行号落在合成语句上),这里会误报失败。改成“包含 expected_line”或至少取该函数最后一帧来比对会更耐重构。

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

已修改为取最后一帧

runtime_value = pto.const(1, dtype=pto.i32)
_ = not runtime_value # AST_REWRITE_TRACEBACK_LINE_MARKER


Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

marker 注释把测试逻辑和源码的物理位置耦合起来了——以后重排/重命名这个文件时要记得同步(唯一性检查能兜底,可以接受)。

更重要的是覆盖缺口:这里只测了模块级函数。真正考验 dedent + increment_lineno 组合的是嵌套/缩进的 jit 函数(getsourcelines 返回带缩进的源码)和带多行装饰器的函数(source_start_line 指向首个 @ 行),建议各补一个用例。

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

修改了测试例,现在的 jit 函数同时是嵌套缩进和多行装饰器的形式

Comment thread ptodsl/tests/test_jit_diagnostics.py Outdated
"ast_rewrite_traceback_line_probe",
"AST_REWRITE_TRACEBACK_LINE_MARKER",
)
rewritten_line_probe = rewrite_jit_function(ast_rewrite_code_line_probe)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这里直接调用了私有函数 rewrite_jit_function,把内部实现钉进了测试契约;更行为导向的做法是通过 @pto.jit 装饰后的产物检查 co_firstlineno。考虑到本文件已有从 _ast_rewrite 导入的先例,不算问题,提一下。

@liggest liggest Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

修改测试例后不再导入 rewrite_jit_function 了,也检查了 co_firstlineno

@liggest
liggest force-pushed the fix_ast_rewrite_tb branch from 131fdde to c30462a Compare August 7, 2026 06:55
@liggest
liggest requested a review from Zhendong404 August 7, 2026 07:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] PTODSL's ast_rewrite reports incorrect traceback line numbers

2 participants