422 lines
16 KiB
Python
422 lines
16 KiB
Python
import t, c
|
|
from stdint import *
|
|
import ast
|
|
import string
|
|
import memhub
|
|
|
|
|
|
# ============================================================
|
|
# StubGen.Converter - pyi 存根文件生成器
|
|
#
|
|
# 从 AST 直接生成 .pyi 存根文件，不依赖 CPython 的 PythonToStubConverter。
|
|
# 遍历 Module 节点的 children，提取函数签名、类定义、变量声明、
|
|
# import 语句，生成 Python 类型存根文件。
|
|
#
|
|
# 本文件使用 Viper 语法实现，可被 TPC 编译为原生码。
|
|
# ============================================================
|
|
|
|
|
|
# ============================================================
|
|
# _PyiEmit - 追加字符串到缓冲区
|
|
#
|
|
# 将 text 追加到 buf 的 pos 位置，返回新的 pos。
|
|
# 自动处理缓冲区溢出（截断）。
|
|
# ============================================================
|
|
def _PyiEmit(buf: bytes, buf_size: t.CSizeT, pos: t.CSizeT,
|
|
text: str) -> t.CSizeT:
|
|
"""追加 text 到 buf 的 pos 位置，返回新的 pos"""
|
|
if buf is None or text is None or pos >= buf_size:
|
|
return pos
|
|
tlen: t.CSizeT = string.strlen(text)
|
|
if tlen == 0:
|
|
return pos
|
|
remain: t.CSizeT = buf_size - pos
|
|
if tlen >= remain:
|
|
tlen = remain - 1
|
|
if tlen > 0:
|
|
string.memcpy(buf + pos, text, tlen)
|
|
return pos + tlen
|
|
|
|
|
|
# ============================================================
|
|
# _PyiIndent - 追加缩进空格
|
|
#
|
|
# 追加 level*4 个空格到 buf 的 pos 位置。
|
|
# ============================================================
|
|
def _PyiIndent(buf: bytes, buf_size: t.CSizeT, pos: t.CSizeT,
|
|
level: int) -> t.CSizeT:
|
|
"""追加 level*4 个空格到 buf，返回新的 pos"""
|
|
i: int = 0
|
|
while i < level:
|
|
pos = _PyiEmit(buf, buf_size, pos, " ")
|
|
i += 1
|
|
return pos
|
|
|
|
|
|
# ============================================================
|
|
# _PyiTypeStr - 递归生成类型注解字符串
|
|
#
|
|
# 将 AST 类型注解节点转换为字符串表示，写入 buf 的 pos 位置。
|
|
# 支持: Name, Attribute, Subscript, BinOp(|), Constant(前向引用), Tuple
|
|
# ============================================================
|
|
def _PyiTypeStr(buf: bytes, buf_size: t.CSizeT, pos: t.CSizeT,
|
|
annot: ast.AST | t.CPtr) -> t.CSizeT:
|
|
"""递归生成类型注解字符串，返回新的 pos"""
|
|
if annot is None:
|
|
return _PyiEmit(buf, buf_size, pos, "t.CInt")
|
|
|
|
k: int = annot.kind()
|
|
|
|
# Name: 简单类型名 (如 int, str, MyClass)
|
|
if k == ast.ASTKind.Name:
|
|
nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(annot)
|
|
if nm.id is not None:
|
|
return _PyiEmit(buf, buf_size, pos, nm.id)
|
|
return _PyiEmit(buf, buf_size, pos, "t.CInt")
|
|
|
|
# Attribute: 属性引用 (如 t.CInt, memhub.MemBuddy)
|
|
if k == ast.ASTKind.Attribute:
|
|
at: ast.Attribute | t.CPtr = (ast.Attribute | t.CPtr)(annot)
|
|
pos = _PyiTypeStr(buf, buf_size, pos, at.value)
|
|
pos = _PyiEmit(buf, buf_size, pos, ".")
|
|
if at.attr is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, at.attr)
|
|
return pos
|
|
|
|
# Subscript: 下标类型 (如 list[int], t.CArray[elem, count])
|
|
if k == ast.ASTKind.Subscript:
|
|
sub: ast.Subscript | t.CPtr = (ast.Subscript | t.CPtr)(annot)
|
|
pos = _PyiTypeStr(buf, buf_size, pos, sub.value)
|
|
pos = _PyiEmit(buf, buf_size, pos, "[")
|
|
pos = _PyiTypeStr(buf, buf_size, pos, sub.slice)
|
|
pos = _PyiEmit(buf, buf_size, pos, "]")
|
|
return pos
|
|
|
|
# BinOp: 联合类型 (如 int | str, File | t.CPtr)
|
|
if k == ast.ASTKind.BinOp:
|
|
bop: ast.BinOp | t.CPtr = (ast.BinOp | t.CPtr)(annot)
|
|
pos = _PyiTypeStr(buf, buf_size, pos, bop.left)
|
|
pos = _PyiEmit(buf, buf_size, pos, " | ")
|
|
pos = _PyiTypeStr(buf, buf_size, pos, bop.right)
|
|
return pos
|
|
|
|
# Constant: 字符串前向引用 (如 'MyClass') 或 None
|
|
if k == ast.ASTKind.Constant:
|
|
cn: ast.Constant | t.CPtr = (ast.Constant | t.CPtr)(annot)
|
|
if cn.const_kind == ast.CONST_STR:
|
|
pos = _PyiEmit(buf, buf_size, pos, "'")
|
|
if cn.str_val is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, cn.str_val)
|
|
pos = _PyiEmit(buf, buf_size, pos, "'")
|
|
return pos
|
|
if cn.const_kind == ast.CONST_NONE:
|
|
return _PyiEmit(buf, buf_size, pos, "None")
|
|
return _PyiEmit(buf, buf_size, pos, "t.CInt")
|
|
|
|
# Tuple: 多参数下标 (如 Dict[str, int])
|
|
if k == ast.ASTKind.Tuple:
|
|
tup: ast.Tuple | t.CPtr = (ast.Tuple | t.CPtr)(annot)
|
|
if tup.elts is not None:
|
|
tn: t.CSizeT = tup.elts.__len__()
|
|
ti: t.CSizeT = 0
|
|
while ti < tn:
|
|
if ti > 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, ", ")
|
|
el: ast.AST | t.CPtr = tup.elts.get(ti)
|
|
pos = _PyiTypeStr(buf, buf_size, pos, el)
|
|
ti += 1
|
|
return pos
|
|
|
|
# 未知类型，回退到 t.CInt
|
|
return _PyiEmit(buf, buf_size, pos, "t.CInt")
|
|
|
|
|
|
# ============================================================
|
|
# _PyiFunction - 生成函数签名
|
|
#
|
|
# 格式: def name(param: Type, ...) -> ReturnType: pass
|
|
# ============================================================
|
|
def _PyiFunction(buf: bytes, buf_size: t.CSizeT, pos: t.CSizeT,
|
|
fn: ast.FunctionDef | t.CPtr,
|
|
indent_level: int, class_name: str) -> t.CSizeT:
|
|
"""生成函数签名到 buf，返回新的 pos"""
|
|
pos = _PyiIndent(buf, buf_size, pos, indent_level)
|
|
pos = _PyiEmit(buf, buf_size, pos, "def ")
|
|
if fn.name is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, fn.name)
|
|
pos = _PyiEmit(buf, buf_size, pos, "(")
|
|
|
|
# 参数列表
|
|
args_node: ast.AST | t.CPtr = fn.args
|
|
if args_node is not None:
|
|
args: ast.Arguments | t.CPtr = (ast.Arguments | t.CPtr)(args_node)
|
|
first: int = 1
|
|
|
|
# 位置参数
|
|
if args.args is not None:
|
|
an: t.CSizeT = args.args.__len__()
|
|
ai: t.CSizeT = 0
|
|
while ai < an:
|
|
if first == 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, ", ")
|
|
first = 0
|
|
ag: ast.Arg | t.CPtr = (ast.Arg | t.CPtr)(args.args.get(ai))
|
|
if ag.arg is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, ag.arg)
|
|
if ag.annotation is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, ": ")
|
|
pos = _PyiTypeStr(buf, buf_size, pos, ag.annotation)
|
|
ai += 1
|
|
|
|
# *args
|
|
if args.vararg is not None:
|
|
if first == 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, ", ")
|
|
first = 0
|
|
pos = _PyiEmit(buf, buf_size, pos, "*")
|
|
va: ast.Arg | t.CPtr = (ast.Arg | t.CPtr)(args.vararg)
|
|
if va.arg is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, va.arg)
|
|
|
|
# **kwargs
|
|
if args.kwarg is not None:
|
|
if first == 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, ", ")
|
|
first = 0
|
|
pos = _PyiEmit(buf, buf_size, pos, "**")
|
|
kw: ast.Arg | t.CPtr = (ast.Arg | t.CPtr)(args.kwarg)
|
|
if kw.arg is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, kw.arg)
|
|
|
|
pos = _PyiEmit(buf, buf_size, pos, ")")
|
|
|
|
# 返回类型
|
|
if fn.returns is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, " -> ")
|
|
pos = _PyiTypeStr(buf, buf_size, pos, fn.returns)
|
|
|
|
pos = _PyiEmit(buf, buf_size, pos, ": pass\n")
|
|
return pos
|
|
|
|
|
|
# ============================================================
|
|
# _PyiClass - 生成类定义
|
|
#
|
|
# 格式:
|
|
# class Name(Base):
|
|
# field: Type
|
|
# def method(self: Name, ...) -> RetType: pass
|
|
# ============================================================
|
|
def _PyiClass(buf: bytes, buf_size: t.CSizeT, pos: t.CSizeT,
|
|
cls: ast.ClassDef | t.CPtr, indent_level: int) -> t.CSizeT:
|
|
"""生成类定义到 buf，返回新的 pos"""
|
|
pos = _PyiIndent(buf, buf_size, pos, indent_level)
|
|
pos = _PyiEmit(buf, buf_size, pos, "class ")
|
|
if cls.name is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, cls.name)
|
|
|
|
# 基类
|
|
if cls.bases is not None:
|
|
bn: t.CSizeT = cls.bases.__len__()
|
|
if bn > 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, "(")
|
|
bi: t.CSizeT = 0
|
|
while bi < bn:
|
|
if bi > 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, ", ")
|
|
base: ast.AST | t.CPtr = cls.bases.get(bi)
|
|
pos = _PyiTypeStr(buf, buf_size, pos, base)
|
|
bi += 1
|
|
pos = _PyiEmit(buf, buf_size, pos, ")")
|
|
|
|
pos = _PyiEmit(buf, buf_size, pos, ":\n")
|
|
|
|
# 类成员
|
|
has_member: int = 0
|
|
|
|
if cls.children is not None:
|
|
cn: t.CSizeT = cls.children.__len__()
|
|
ci: t.CSizeT = 0
|
|
while ci < cn:
|
|
stmt: ast.AST | t.CPtr = cls.children.get(ci)
|
|
if stmt is not None:
|
|
sk: int = stmt.kind()
|
|
|
|
# 方法
|
|
if sk == ast.ASTKind.FunctionDef:
|
|
mfn: ast.FunctionDef | t.CPtr = (ast.FunctionDef | t.CPtr)(stmt)
|
|
pos = _PyiFunction(buf, buf_size, pos, mfn, indent_level + 1, cls.name)
|
|
has_member = 1
|
|
|
|
# 带注解的属性
|
|
elif sk == ast.ASTKind.AnnAssign:
|
|
aa: ast.AnnAssign | t.CPtr = (ast.AnnAssign | t.CPtr)(stmt)
|
|
if aa.target is not None and aa.target.kind() == ast.ASTKind.Name:
|
|
tn: ast.Name | t.CPtr = (ast.Name | t.CPtr)(aa.target)
|
|
pos = _PyiIndent(buf, buf_size, pos, indent_level + 1)
|
|
if tn.id is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, tn.id)
|
|
pos = _PyiEmit(buf, buf_size, pos, ": ")
|
|
pos = _PyiTypeStr(buf, buf_size, pos, aa.annotation)
|
|
pos = _PyiEmit(buf, buf_size, pos, "\n")
|
|
has_member = 1
|
|
ci += 1
|
|
|
|
# 空类补 pass
|
|
if has_member == 0:
|
|
pos = _PyiIndent(buf, buf_size, pos, indent_level + 1)
|
|
pos = _PyiEmit(buf, buf_size, pos, "pass\n")
|
|
|
|
pos = _PyiEmit(buf, buf_size, pos, "\n")
|
|
return pos
|
|
|
|
|
|
# ============================================================
|
|
# _GeneratePyiFromAst - 从 AST 生成 .pyi 存根文件
|
|
#
|
|
# 遍历 AST Module 节点，提取函数签名、类定义、变量声明、import 语句，
|
|
# 生成 Python 类型存根文件内容到 buf。
|
|
#
|
|
# Args:
|
|
# mb: 内存池（保留以符合调用约定，实际未使用）
|
|
# tree: AST Module 根节点
|
|
# rel_path: 相对路径（用于模块名，当前未使用）
|
|
# buf: 输出缓冲区
|
|
# buf_size: 缓冲区大小
|
|
#
|
|
# Returns:
|
|
# 写入的字节数（0 表示失败）
|
|
# ============================================================
|
|
def _GeneratePyiFromAst(mb: memhub.MemBuddy | t.CPtr, tree: ast.AST | t.CPtr,
|
|
rel_path: str, buf: bytes,
|
|
buf_size: t.CSizeT) -> t.CSizeT:
|
|
"""从 AST 生成 .pyi 存根文件内容到 buf"""
|
|
if tree is None or buf is None or buf_size == 0:
|
|
return 0
|
|
|
|
pos: t.CSizeT = 0
|
|
|
|
# 文件头
|
|
pos = _PyiEmit(buf, buf_size, pos, "\"\"\"Auto-generated Python stub file\"\"\"\n\n")
|
|
|
|
if tree.children is None:
|
|
if pos < buf_size:
|
|
buf[pos] = 0
|
|
return pos
|
|
|
|
n: t.CSizeT = tree.children.__len__()
|
|
|
|
# 第一遍扫描：检测 import t/c
|
|
has_import_t: int = 0
|
|
has_import_c: int = 0
|
|
i: t.CSizeT = 0
|
|
while i < n:
|
|
stmt: ast.AST | t.CPtr = tree.children.get(i)
|
|
if stmt is not None and stmt.kind() == ast.ASTKind.Import:
|
|
imp: ast.Import | t.CPtr = (ast.Import | t.CPtr)(stmt)
|
|
if imp.names is not None:
|
|
an: t.CSizeT = imp.names.__len__()
|
|
ai: t.CSizeT = 0
|
|
while ai < an:
|
|
al: ast.Alias | t.CPtr = (ast.Alias | t.CPtr)(imp.names.get(ai))
|
|
if al is not None and al.name is not None:
|
|
if string.strcmp(al.name, "t") == 0:
|
|
has_import_t = 1
|
|
elif string.strcmp(al.name, "c") == 0:
|
|
has_import_c = 1
|
|
ai += 1
|
|
i += 1
|
|
|
|
# 补充缺失的 import t/c
|
|
if has_import_t == 0 and has_import_c == 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, "import t, c\n\n")
|
|
elif has_import_t == 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, "import t\n\n")
|
|
elif has_import_c == 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, "import c\n\n")
|
|
else:
|
|
pos = _PyiEmit(buf, buf_size, pos, "\n")
|
|
|
|
# 第二遍扫描：生成声明
|
|
i = 0
|
|
while i < n:
|
|
stmt = tree.children.get(i)
|
|
if stmt is not None:
|
|
k: int = stmt.kind()
|
|
|
|
# Import 语句
|
|
if k == ast.ASTKind.Import:
|
|
imp2: ast.Import | t.CPtr = (ast.Import | t.CPtr)(stmt)
|
|
pos = _PyiEmit(buf, buf_size, pos, "import ")
|
|
if imp2.names is not None:
|
|
an2: t.CSizeT = imp2.names.__len__()
|
|
ai2: t.CSizeT = 0
|
|
while ai2 < an2:
|
|
if ai2 > 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, ", ")
|
|
al2: ast.Alias | t.CPtr = (ast.Alias | t.CPtr)(imp2.names.get(ai2))
|
|
if al2 is not None and al2.name is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, al2.name)
|
|
if al2.asname is not None and string.strlen(al2.asname) > 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, " as ")
|
|
pos = _PyiEmit(buf, buf_size, pos, al2.asname)
|
|
ai2 += 1
|
|
pos = _PyiEmit(buf, buf_size, pos, "\n")
|
|
|
|
# ImportFrom 语句
|
|
elif k == ast.ASTKind.ImportFrom:
|
|
ifr: ast.ImportFrom | t.CPtr = (ast.ImportFrom | t.CPtr)(stmt)
|
|
pos = _PyiEmit(buf, buf_size, pos, "from ")
|
|
lvl: int = ifr.level
|
|
while lvl > 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, ".")
|
|
lvl -= 1
|
|
if ifr.module is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, ifr.module)
|
|
pos = _PyiEmit(buf, buf_size, pos, " import ")
|
|
if ifr.names is not None:
|
|
an3: t.CSizeT = ifr.names.__len__()
|
|
ai3: t.CSizeT = 0
|
|
while ai3 < an3:
|
|
if ai3 > 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, ", ")
|
|
al3: ast.Alias | t.CPtr = (ast.Alias | t.CPtr)(ifr.names.get(ai3))
|
|
if al3 is not None and al3.name is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, al3.name)
|
|
if al3.asname is not None and string.strlen(al3.asname) > 0:
|
|
pos = _PyiEmit(buf, buf_size, pos, " as ")
|
|
pos = _PyiEmit(buf, buf_size, pos, al3.asname)
|
|
ai3 += 1
|
|
pos = _PyiEmit(buf, buf_size, pos, "\n")
|
|
|
|
# 函数定义
|
|
elif k == ast.ASTKind.FunctionDef:
|
|
fn: ast.FunctionDef | t.CPtr = (ast.FunctionDef | t.CPtr)(stmt)
|
|
pos = _PyiFunction(buf, buf_size, pos, fn, 0, "")
|
|
pos = _PyiEmit(buf, buf_size, pos, "\n")
|
|
|
|
# 类定义
|
|
elif k == ast.ASTKind.ClassDef:
|
|
cls: ast.ClassDef | t.CPtr = (ast.ClassDef | t.CPtr)(stmt)
|
|
pos = _PyiClass(buf, buf_size, pos, cls, 0)
|
|
|
|
# 带注解的变量声明
|
|
elif k == ast.ASTKind.AnnAssign:
|
|
aa: ast.AnnAssign | t.CPtr = (ast.AnnAssign | t.CPtr)(stmt)
|
|
if aa.target is not None and aa.target.kind() == ast.ASTKind.Name:
|
|
nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(aa.target)
|
|
if nm.id is not None:
|
|
pos = _PyiEmit(buf, buf_size, pos, nm.id)
|
|
pos = _PyiEmit(buf, buf_size, pos, ": ")
|
|
pos = _PyiTypeStr(buf, buf_size, pos, aa.annotation)
|
|
pos = _PyiEmit(buf, buf_size, pos, "\n")
|
|
|
|
i += 1
|
|
|
|
# NUL 终止
|
|
if pos < buf_size:
|
|
buf[pos] = 0
|
|
else:
|
|
buf[buf_size - 1] = 0
|
|
return pos
|