238 lines
10 KiB
Python
238 lines
10 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.includes import t
|
||
from lib.core.SymbolUtils import AnnotationContainsTType, CheckAnnotationHasCInline
|
||
|
||
|
||
class ASTSymbolExtractor:
|
||
"""从 Python 源文件的 AST 中提取所有符号信息"""
|
||
|
||
def __init__(self, symbol_table: SymbolTable):
|
||
self._symtab = symbol_table
|
||
|
||
def extract(self, file_path: str) -> ModuleSymbols:
|
||
"""提取文件中的所有符号信息"""
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
code = f.read()
|
||
tree = ast.parse(code)
|
||
|
||
result = ModuleSymbols(file_path=file_path)
|
||
|
||
# 第一步：收集所有类的成员和匿名类型
|
||
class_members = {}
|
||
for node in tree.body:
|
||
if isinstance(node, ast.ClassDef):
|
||
members = {}
|
||
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 = 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, members, anonymous_types, prefix=''):
|
||
"""递归收集类成员，包括嵌套匿名类型"""
|
||
for item in class_node.body:
|
||
if isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
|
||
member_name = item.target.id
|
||
full_name = f'{prefix}{member_name}' if prefix else member_name
|
||
member_type = CTypeInfo.FromNode(item.annotation, self._symtab)
|
||
if member_type is None:
|
||
member_type = CTypeInfo()
|
||
member_type.BaseType = t.CInt()
|
||
is_ptr = member_type.IsPtr
|
||
if isinstance(item.annotation, ast.BinOp):
|
||
annotation_str = ast.dump(item.annotation)
|
||
if 'CPtr' in annotation_str:
|
||
is_ptr = True
|
||
members[full_name] = member_type.Copy()
|
||
|
||
elif isinstance(item, ast.ClassDef):
|
||
is_anonymous = any(
|
||
any(AnnotationContainsTType(base, 'Anonymous') for base in item.bases)
|
||
) if item.bases else False
|
||
|
||
if is_anonymous:
|
||
is_union = any(
|
||
any(AnnotationContainsTType(base, 'CUnion') for base in item.bases)
|
||
)
|
||
|
||
nested_name = 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 = f'{prefix}{item.name}.' if prefix else f'{item.name}.'
|
||
self._collect_members(item, members, anonymous_types, new_prefix)
|
||
|
||
anonymous_members = {}
|
||
for full_name, member_info in members.items():
|
||
if full_name.startswith(f"{nested_name}."):
|
||
member_name = 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, members_info) -> ClassSymbolData:
|
||
"""从 ClassDef 节点提取类符号数据"""
|
||
type_kind = 'struct'
|
||
is_cpython_object = False
|
||
is_packed = 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
|
||
decorator.value.id == 't' and
|
||
decorator.attr == 'Object'):
|
||
is_cpython_object = True
|
||
break
|
||
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 arg.value.value.id == 't' 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, '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 = self._symtab.get(base.id)
|
||
if base_entry and base_entry.IsExceptionClass:
|
||
type_kind = 'exception'
|
||
break
|
||
|
||
# 提取枚举成员
|
||
enum_members = []
|
||
if type_kind == 'enum':
|
||
next_enum_value = 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 = 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 = 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))
|
||
|
||
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, result: ModuleSymbols):
|
||
"""从 AnnAssign 节点提取 typedef 或 define 数据"""
|
||
var_name = node.target.id
|
||
has_cdefine = AnnotationContainsTType(node.annotation, 'CDefine') if node.annotation else False
|
||
has_postdef = AnnotationContainsTType(node.annotation, 'Postdefinition') if node.annotation else False
|
||
has_ctypedef = 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) -> FuncSymbolData:
|
||
"""从 FunctionDef 节点提取函数符号数据"""
|
||
params = [(arg.arg, arg.annotation) for arg in node.args.args]
|
||
is_variadic = node.args.vararg is not None
|
||
is_inline = 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,
|
||
)
|