"""Pass 1: ASTSymbolExtractor — 从 AST 提取原始符号信息""" from __future__ import annotations import ast from typing import TYPE_CHECKING if TYPE_CHECKING: from lib.core.SymbolTable import SymbolTable from lib.core.SymbolData import ( ModuleSymbols, ClassSymbolData, TypedefSymbolData, FuncSymbolData, DefineSymbolData, AnonymousTypeData ) from lib.core.Handles.HandlesBase import CTypeInfo from lib.core.SymbolUtils import AnnotationContainsTType, CheckAnnotationHasCInline, IsTModule, AnnotationContainsName from lib.includes import t class ASTSymbolExtractor: """从 Python 源文件的 AST 中提取所有符号信息""" def __init__(self, symbol_table: SymbolTable) -> None: self._symtab: SymbolTable = symbol_table def extract(self, file_path: str) -> ModuleSymbols: """提取文件中的所有符号信息""" with open(file_path, 'r', encoding='utf-8') as f: code: str = f.read() tree: ast.Module = ast.parse(code) result: ModuleSymbols = ModuleSymbols(file_path=file_path) # 第一步：收集所有类的成员和匿名类型 class_members: dict[str, dict[str, CTypeInfo]] = {} for node in tree.body: if isinstance(node, ast.ClassDef): members: dict[str, CTypeInfo] = {} self._collect_members(node, members, result.anonymous_types) class_members[node.name] = members # 第二步：从每个顶层节点提取符号信息 for node in tree.body: if isinstance(node, ast.ClassDef): members_info: dict[str, CTypeInfo] = class_members.get(node.name, {}) result.classes.append(self._extract_class(node, members_info)) elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): self._extract_ann_assign(node, result) elif isinstance(node, ast.FunctionDef): result.functions.append(self._extract_function(node)) return result # ------------------------------------------------------------------ # 成员收集（递归处理嵌套匿名类） # ------------------------------------------------------------------ def _collect_members(self, class_node: ast.ClassDef, members: dict[str, CTypeInfo], anonymous_types: dict[str, AnonymousTypeData], prefix: str = '') -> None: """递归收集类成员，包括嵌套匿名类型""" for item in class_node.body: if isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name): member_name: str = item.target.id full_name: str = f'{prefix}{member_name}' if prefix else member_name member_type: CTypeInfo | None = CTypeInfo.FromNode(item.annotation, self._symtab) if member_type is None: member_type = CTypeInfo() member_type.BaseType = t.CInt() is_ptr: bool = member_type.IsPtr if isinstance(item.annotation, ast.BinOp): if AnnotationContainsName(item.annotation, 'CPtr'): is_ptr = True members[full_name] = member_type.Copy() elif isinstance(item, ast.ClassDef): is_anonymous: bool = any( any(AnnotationContainsTType(base, 'Anonymous') for base in item.bases) ) if item.bases else False if is_anonymous: is_union: bool = any( any(AnnotationContainsTType(base, 'CUnion') for base in item.bases) ) # 也检查装饰器形式 @t.CUnion if not is_union and hasattr(item, 'decorator_list') and item.decorator_list: for deco in item.decorator_list: if isinstance(deco, ast.Attribute) and hasattr(deco.value, 'id') and IsTModule(deco.value.id, self._symtab) and deco.attr == 'CUnion': is_union = True break elif isinstance(deco, ast.Name): cls2: type | None = getattr(t, deco.id, None) if cls2 is None and hasattr(self._symtab, '_t_type_symbols'): cls2 = self._symtab._t_type_symbols.get(deco.id) if cls2 is not None and isinstance(cls2, type) and issubclass(cls2, t.CUnion): is_union = True break nested_name: str = f'{prefix}{item.name}' if prefix else item.name members[nested_name] = CTypeInfo() members[nested_name].BaseType = f'union {item.name}' if is_union else f'struct {item.name}' members[nested_name].PtrCount = 0 members[nested_name].ArrayDims = [] new_prefix: str = f'{prefix}{item.name}.' if prefix else f'{item.name}.' self._collect_members(item, members, anonymous_types, new_prefix) anonymous_members: dict[str, CTypeInfo] = {} for full_name, member_info in members.items(): if full_name.startswith(f"{nested_name}."): member_name: str = full_name[len(nested_name) + 1:] anonymous_members[member_name] = member_info if item.name not in anonymous_types: anonymous_types[item.name] = AnonymousTypeData( name=item.name, is_union=is_union, members=anonymous_members, lineno=item.lineno ) # ------------------------------------------------------------------ # 类符号提取 # ------------------------------------------------------------------ def _extract_class(self, node: ast.ClassDef, members_info: dict[str, CTypeInfo]) -> ClassSymbolData: """从 ClassDef 节点提取类符号数据""" type_kind: str = 'struct' is_cpython_object: bool = False is_packed: bool = False # 检查装饰器 if hasattr(node, 'decorator_list') and node.decorator_list: for decorator in node.decorator_list: if isinstance(decorator, ast.Attribute): if (hasattr(decorator.value, 'id') and IsTModule(decorator.value.id, self._symtab)): if decorator.attr == 'Object': is_cpython_object = True elif decorator.attr == 'CUnion': type_kind = 'union' elif decorator.attr == 'CEnum': type_kind = 'enum' elif decorator.attr == 'REnum': type_kind = 'renum' elif decorator.attr == 'CStruct': type_kind = 'struct' elif isinstance(decorator, ast.Name): # from t import CStruct / CUnion / CEnum / REnum cls: type | None = getattr(t, decorator.id, None) if cls is None and hasattr(self._symtab, '_t_type_symbols'): cls = self._symtab._t_type_symbols.get(decorator.id) if cls is not None and isinstance(cls, type) and issubclass(cls, t.CType): if issubclass(cls, t.CUnion): type_kind = 'union' elif issubclass(cls, t.REnum): type_kind = 'renum' elif issubclass(cls, t.CEnum): type_kind = 'enum' elif issubclass(cls, t.CStruct): type_kind = 'struct' elif decorator.id == 'Object': is_cpython_object = True elif isinstance(decorator, ast.Call): if isinstance(decorator.func, ast.Attribute): if hasattr(decorator.func.value, 'id') and decorator.func.value.id == 'c' and decorator.func.attr == 'Attribute': for arg in decorator.args: if isinstance(arg, ast.Attribute): if isinstance(arg.value, ast.Attribute): if hasattr(arg.value.value, 'id') and IsTModule(arg.value.value.id, self._symtab) and arg.value.attr == 'attr' and arg.attr == 'packed': is_packed = True # 检查基类 for base in node.bases: if AnnotationContainsTType(base, 'CUnion'): type_kind = 'union' break elif AnnotationContainsTType(base, 'CEnum'): type_kind = 'enum' break elif AnnotationContainsTType(base, 'REnum'): type_kind = 'renum' break elif AnnotationContainsTType(base, 'CStruct'): type_kind = 'struct' break elif AnnotationContainsTType(base, 'Object'): is_cpython_object = True type_kind = 'struct' break elif isinstance(base, ast.Name) and base.id == 'Exception': type_kind = 'exception' break elif isinstance(base, ast.Name): base_entry: CTypeInfo | None = self._symtab.get(base.id) if base_entry and base_entry.IsExceptionClass: type_kind = 'exception' break # 提取枚举成员 enum_members: list[tuple[str, int, int]] = [] if type_kind == 'enum': next_enum_value: int = 0 for item in node.body: if isinstance(item, ast.Assign): if len(item.targets) == 1 and isinstance(item.targets[0], ast.Name): member_name: str = item.targets[0].id if isinstance(item.value, ast.Constant) and isinstance(item.value.value, int): next_enum_value = item.value.value + 1 else: next_enum_value += 1 enum_members.append((member_name, next_enum_value - 1, item.lineno)) elif isinstance(item, ast.AnnAssign): if isinstance(item.target, ast.Name): member_name: str = item.target.id if item.value: if isinstance(item.value, ast.Constant) and isinstance(item.value.value, int): next_enum_value = item.value.value + 1 else: next_enum_value += 1 else: next_enum_value += 1 enum_members.append((member_name, next_enum_value - 1, item.lineno)) # REnum 变体提取（嵌套 ClassDef + Assign 显式值） # 与 HandlesClassDef._EmitREnumLlvm 的 tag 分配逻辑对齐： # - 先收集 Assign 显式值 # - ClassDef 变体从 max(explicit)+1 开始顺序编号 # - Assign 无显式值时用独立计数器 variant_index if type_kind == 'renum': explicit_values: dict[str, int] = {} for item in node.body: if isinstance(item, ast.Assign): for target in item.targets: if isinstance(target, ast.Name): if item.value and isinstance(item.value, ast.Constant) and isinstance(item.value.value, int): explicit_values[target.id] = item.value.value next_class_tag: int = max(explicit_values.values()) + 1 if explicit_values else 0 variant_index: int = 0 for item in node.body: if isinstance(item, ast.ClassDef): variant_name_r: str = item.name tag: int = next_class_tag next_class_tag += 1 enum_members.append((variant_name_r, tag, item.lineno)) elif isinstance(item, ast.Assign): for target in item.targets: if isinstance(target, ast.Name): var_name: str = target.id if item.value and isinstance(item.value, ast.Constant) and isinstance(item.value.value, int): value: int = item.value.value else: value = variant_index variant_index += 1 enum_members.append((var_name, value, item.lineno)) return ClassSymbolData( name=node.name, type_kind=type_kind, lineno=node.lineno, members=members_info, is_cpython_object=is_cpython_object, is_packed=is_packed, enum_members=enum_members, ) # ------------------------------------------------------------------ # AnnAssign 提取（typedef / define） # ------------------------------------------------------------------ def _extract_ann_assign(self, node: ast.AnnAssign, result: ModuleSymbols) -> None: """从 AnnAssign 节点提取 typedef 或 define 数据""" var_name: str = node.target.id has_cdefine: bool = AnnotationContainsTType(node.annotation, 'CDefine') if node.annotation else False has_postdef: bool = AnnotationContainsTType(node.annotation, 'Postdefinition') if node.annotation else False has_ctypedef: bool = AnnotationContainsTType(node.annotation, 'CTypedef') if node.annotation else False if has_cdefine: result.defines.append(DefineSymbolData( name=var_name, lineno=node.lineno, value_node=node.value, )) else: result.typedefs.append(TypedefSymbolData( name=var_name, lineno=node.lineno, annotation=node.annotation, value=node.value, has_cdefine=has_cdefine, has_postdef=has_postdef, has_ctypedef=has_ctypedef, )) # ------------------------------------------------------------------ # 函数提取 # ------------------------------------------------------------------ def _extract_function(self, node: ast.FunctionDef) -> FuncSymbolData: """从 FunctionDef 节点提取函数符号数据""" params: list[tuple[str, ast.AST | None]] = [(arg.arg, arg.annotation) for arg in node.args.args] is_variadic: bool = node.args.vararg is not None is_inline: bool = CheckAnnotationHasCInline(node.returns) return FuncSymbolData( name=node.name, lineno=node.lineno, returns_node=node.returns, params=params, is_variadic=is_variadic, is_inline=is_inline, )