308 lines
15 KiB
Python
308 lines
15 KiB
Python
"""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,
|
||
)
|