TPC 侧取缔对 Python 标准库 ast 模块的依赖,改用 lib/vast 包: - Nodes: 与 CPython 3.12 ast 完全同构的 131 个节点类(含抽象基类继承) - Tokenizer: 词法分析(缩进/字符串/f-string/数字/运算符/软关键字) - ExprParser: 表达式递归下降解析(优先级/lambda/推导式/f-string/切片) - Parser: 语句解析(def/class/if/for/try/with/match/import/注解) - Utils/Unparser: walk/dump/unparse/copy_location/literal_eval 等辅助 - 软关键字 fn/fun/func/function 并入解析器, 废除 ExpandSoftKeywords 预处理 - lib/_bootstrap 将 lib 目录加入 sys.path, import vast as ast 零修改替换 - 52 个文件 import ast 替换为 import vast as ast; 94 文件差分测试 0 失败
282 lines
11 KiB
Python
282 lines
11 KiB
Python
"""vast.utils - AST 遍历与序列化工具
|
||
|
||
与 CPython ``ast`` 模块辅助函数同构:
|
||
|
||
- ``iter_fields`` / ``iter_child_nodes`` / ``walk``
|
||
- ``dump``(含 annotate_fields / include_attributes / indent)
|
||
- ``copy_location`` / ``fix_missing_locations`` / ``get_source_segment``
|
||
- ``literal_eval``(常量求值)
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from lib.vast.Nodes import AST
|
||
from lib.vast.Tokenizer import Token
|
||
|
||
|
||
def iter_fields(node: AST):
|
||
"""按 _fields 顺序产出 (字段名, 值)。"""
|
||
for field in node._fields:
|
||
try:
|
||
yield field, getattr(node, field)
|
||
except AttributeError:
|
||
pass
|
||
|
||
|
||
def iter_child_nodes(node: AST):
|
||
"""产出全部子节点(递归扁平,跳过 None 与空列表)。"""
|
||
for _name, field in iter_fields(node):
|
||
if isinstance(field, AST):
|
||
yield field
|
||
elif isinstance(field, list):
|
||
for item in field:
|
||
if isinstance(item, AST):
|
||
yield item
|
||
|
||
|
||
def walk(node: AST):
|
||
"""深度优先遍历,产出自身与全部后代节点。"""
|
||
from collections import deque
|
||
|
||
todo: deque[AST] = deque([node])
|
||
while todo:
|
||
n: AST = todo.popleft()
|
||
yield n
|
||
todo.extend(iter_child_nodes(n))
|
||
|
||
|
||
def dump(node: AST, annotate_fields: bool = True, include_attributes: bool = False,
|
||
indent: int | str | None = None) -> str:
|
||
"""将 AST 节点格式化为字符串(精确复刻 CPython 3.12 ast.dump)。
|
||
|
||
- annotate_fields: 是否显示字段名(默认 True)
|
||
- include_attributes: 是否附加 lineno/col_offset 等位置属性
|
||
- indent: 缩进(int 或字符串),None 为单行紧凑输出
|
||
|
||
与 CPython 一致的行为:
|
||
- 值为 None 且类默认值为 None 的字段跳过(如 FunctionDef.returns)
|
||
- 节点字段全部为"简单值"且 ≤3 个时单行输出
|
||
"""
|
||
if not isinstance(node, AST):
|
||
raise TypeError(f'expected AST, got {type(node).__name__!r}')
|
||
|
||
if indent is not None and not isinstance(indent, str):
|
||
indent = ' ' * indent
|
||
|
||
def _format(node: object, level: int = 0):
|
||
if indent is not None:
|
||
level += 1
|
||
prefix: str = '\n' + indent * level # type: ignore[operator]
|
||
sep: str = ',\n' + indent * level # type: ignore[operator]
|
||
else:
|
||
prefix = ''
|
||
sep = ', '
|
||
if isinstance(node, AST):
|
||
cls: type = type(node)
|
||
args: list[str] = []
|
||
allsimple: bool = True
|
||
keywords: bool = annotate_fields
|
||
for name in node._fields:
|
||
try:
|
||
value = getattr(node, name)
|
||
except AttributeError:
|
||
keywords = True
|
||
continue
|
||
if value is None and getattr(cls, name, ...) is None:
|
||
keywords = True
|
||
continue
|
||
value, simple = _format(value, level)
|
||
allsimple = allsimple and simple
|
||
if keywords:
|
||
args.append(f'{name}={value}')
|
||
else:
|
||
args.append(value)
|
||
if include_attributes and node._attributes:
|
||
for name in node._attributes:
|
||
try:
|
||
value = getattr(node, name)
|
||
except AttributeError:
|
||
continue
|
||
if value is None and getattr(cls, name, ...) is None:
|
||
continue
|
||
value, simple = _format(value, level)
|
||
allsimple = allsimple and simple
|
||
args.append(f'{name}={value}')
|
||
if allsimple and len(args) <= 3:
|
||
return f'{type(node).__name__}({", ".join(args)})', not args
|
||
return f'{type(node).__name__}({prefix}{sep.join(args)})', False
|
||
elif isinstance(node, list):
|
||
if not node:
|
||
return '[]', True
|
||
return f'[{prefix}{sep.join(_format(x, level)[0] for x in node)}]', False
|
||
return repr(node), True
|
||
|
||
return _format(node)[0]
|
||
|
||
|
||
def copy_location(new_node: AST, old_node: AST) -> AST:
|
||
"""复制 old_node 的位置属性到 new_node(与 ast.copy_location 一致)。"""
|
||
for attr in ('lineno', 'col_offset', 'end_lineno', 'end_col_offset'):
|
||
try:
|
||
setattr(new_node, attr, getattr(old_node, attr))
|
||
except AttributeError:
|
||
pass
|
||
return new_node
|
||
|
||
|
||
def fix_missing_locations(node: AST) -> AST:
|
||
"""为缺失位置的节点补全(从最近的祖先复制;递归处理)。"""
|
||
def _fix(n: AST, parent_loc: dict[str, int] | None) -> None:
|
||
has_loc: bool = hasattr(n, 'lineno') and getattr(n, 'lineno', None) is not None
|
||
if hasattr(n, 'lineno') and getattr(n, 'lineno', None) is None and parent_loc:
|
||
for attr, val in parent_loc.items():
|
||
setattr(n, attr, val)
|
||
loc: dict[str, int] | None = None
|
||
if hasattr(n, 'lineno') and getattr(n, 'lineno', None) is not None:
|
||
loc = {a: getattr(n, a) for a in
|
||
('lineno', 'col_offset', 'end_lineno', 'end_col_offset')
|
||
if getattr(n, a, None) is not None}
|
||
for child in iter_child_nodes(n):
|
||
_fix(child, loc)
|
||
|
||
_fix(node, None)
|
||
return node
|
||
|
||
|
||
def get_source_segment(source: str, node: AST, padded: bool = False) -> str | None:
|
||
"""提取节点对应的源码片段(与 ast.get_source_segment 一致)。"""
|
||
try:
|
||
lineno: int | None = getattr(node, 'lineno', None)
|
||
end_lineno: int | None = getattr(node, 'end_lineno', None)
|
||
col_offset: int | None = getattr(node, 'col_offset', None)
|
||
end_col_offset: int | None = getattr(node, 'end_col_offset', None)
|
||
if lineno is None or end_lineno is None or col_offset is None or end_col_offset is None:
|
||
return None
|
||
lines: list[str] = source.splitlines(keepends=True)
|
||
if lineno < 1 or lineno > len(lines):
|
||
return None
|
||
first: str = lines[lineno - 1]
|
||
# col_offset 是 UTF-8 字节偏移 → 转字符偏移
|
||
start_char: int = _ByteToChar(first, col_offset)
|
||
if lineno == end_lineno:
|
||
end_char: int = _ByteToChar(first, end_col_offset)
|
||
segment: str = first[start_char:end_char]
|
||
else:
|
||
last: str = lines[end_lineno - 1]
|
||
end_char = _ByteToChar(last, end_col_offset)
|
||
middle: list[str] = lines[lineno:end_lineno - 1]
|
||
segment = first[start_char:] + ''.join(middle) + last[:end_char]
|
||
if padded:
|
||
return segment
|
||
return segment
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _ByteToChar(line: str, byte_offset: int) -> int:
|
||
"""UTF-8 字节偏移 → 字符偏移。"""
|
||
b: bytes = line.encode('utf-8')
|
||
return len(b[:byte_offset].decode('utf-8', errors='replace')) if byte_offset > 0 else 0
|
||
|
||
|
||
def literal_eval(node_or_string: AST | str):
|
||
"""安全求值常量表达式(与 ast.literal_eval 语义一致)。
|
||
|
||
支持常量、元组/列表/集合/字典、一元 +/-、BinOp 算术、切片。
|
||
"""
|
||
if isinstance(node_or_string, str):
|
||
from lib.vast.Tokenizer import Tokenizer
|
||
|
||
toks: list[Token] = Tokenizer(node_or_string).Tokenize()
|
||
if toks and toks[0].Type == 'NEWLINE':
|
||
raise ValueError('malformed node or string')
|
||
from lib.vast.Parser import _ComputeLineStarts
|
||
|
||
from lib.vast.ExprParser import ExprParserMixin
|
||
|
||
class _EV(ExprParserMixin):
|
||
def __init__(self) -> None:
|
||
super().__init__(toks, node_or_string, _ComputeLineStarts(node_or_string))
|
||
|
||
parser = _EV()
|
||
node: AST = parser.ParseExpr()
|
||
elif isinstance(node_or_string, AST):
|
||
node = node_or_string
|
||
else:
|
||
raise TypeError('bad argument type')
|
||
|
||
return _EvalLiteral(node)
|
||
|
||
|
||
def _EvalLiteral(node: AST):
|
||
"""递归常量求值核心。"""
|
||
if isinstance(node, AST):
|
||
pass
|
||
name: str = type(node).__name__
|
||
# 常量
|
||
if name == 'Constant':
|
||
value = node.value
|
||
if isinstance(value, (int, float, complex, str, bool, bytes, type(None), Ellipsis)): # type: ignore[comparison-overlap]
|
||
return value
|
||
raise ValueError('malformed node or string: unsupported constant')
|
||
if name == 'Tuple':
|
||
return tuple(_EvalLiteral(e) for e in node.elts)
|
||
if name == 'List':
|
||
return [_EvalLiteral(e) for e in node.elts]
|
||
if name == 'Set':
|
||
return {_EvalLiteral(e) for e in node.elts}
|
||
if name == 'Dict':
|
||
if len(node.keys) != len(node.values):
|
||
raise ValueError('malformed node or string: Dict')
|
||
out: dict = {}
|
||
for k, v in zip(node.keys, node.values):
|
||
if k is None:
|
||
raise ValueError('malformed node or string: ** in Dict')
|
||
out[_EvalLiteral(k)] = _EvalLiteral(v)
|
||
return out
|
||
if name == 'UnaryOp':
|
||
operand = _EvalLiteral(node.operand)
|
||
if isinstance(node.op, (__import__('lib.vast.Nodes', fromlist=['USub']).USub,)):
|
||
return -operand
|
||
if isinstance(node.op, (__import__('lib.vast.Nodes', fromlist=['UAdd']).UAdd,)):
|
||
return +operand
|
||
if isinstance(node.op, (__import__('lib.vast.Nodes', fromlist=['Not']).Not,)):
|
||
return not operand
|
||
raise ValueError('malformed node or string: unary op')
|
||
if name == 'BinOp':
|
||
from lib.vast.Nodes import Add, Sub, Mult, Div, FloorDiv, Mod, Pow, LShift, RShift, BitOr, BitXor, BitAnd
|
||
|
||
left = _EvalLiteral(node.left)
|
||
right = _EvalLiteral(node.right)
|
||
op = node.op
|
||
if isinstance(op, Add):
|
||
return left + right
|
||
if isinstance(op, Sub):
|
||
return left - right
|
||
if isinstance(op, Mult):
|
||
return left * right
|
||
if isinstance(op, Div):
|
||
return left / right
|
||
if isinstance(op, FloorDiv):
|
||
return left // right
|
||
if isinstance(op, Mod):
|
||
return left % right
|
||
if isinstance(op, Pow):
|
||
return left ** right
|
||
if isinstance(op, LShift):
|
||
return left << right
|
||
if isinstance(op, RShift):
|
||
return left >> right
|
||
if isinstance(op, BitOr):
|
||
return left | right
|
||
if isinstance(op, BitXor):
|
||
return left ^ right
|
||
if isinstance(op, BitAnd):
|
||
return left & right
|
||
raise ValueError('malformed node or string: bin op')
|
||
if name == 'Slice':
|
||
lower = _EvalLiteral(node.lower) if node.lower is not None else None
|
||
upper = _EvalLiteral(node.upper) if node.upper is not None else None
|
||
step = _EvalLiteral(node.step) if node.step is not None else None
|
||
return slice(lower, upper, step)
|
||
raise ValueError(f'malformed node or string: {name}')
|