|
|
""" |
|
|
This generates .pyi stubs for the cffi Python bindings generated by regenerate.py |
|
|
""" |
|
|
import sys, re, itertools |
|
|
sys.path.extend(['.', '..']) |
|
|
|
|
|
from pycparser import c_ast, parse_file, CParser |
|
|
import pycparser.plyparser |
|
|
from pycparser.c_ast import PtrDecl, TypeDecl, FuncDecl, EllipsisParam, IdentifierType, Struct, Enum, Typedef |
|
|
from typing import Tuple |
|
|
|
|
|
__c_type_to_python_type = { |
|
|
'void': 'None', '_Bool': 'bool', |
|
|
'char': 'int', 'short': 'int', 'int': 'int', 'long': 'int', |
|
|
'ptrdiff_t': 'int', 'size_t': 'int', |
|
|
'int8_t': 'int', 'uint8_t': 'int', |
|
|
'int16_t': 'int', 'uint16_t': 'int', |
|
|
'int32_t': 'int', 'uint32_t': 'int', |
|
|
'int64_t': 'int', 'uint64_t': 'int', |
|
|
'float': 'float', 'double': 'float', |
|
|
'ggml_fp16_t': 'np.float16', |
|
|
} |
|
|
|
|
|
def format_type(t: TypeDecl): |
|
|
if isinstance(t, PtrDecl) or isinstance(t, Struct): |
|
|
return 'ffi.CData' |
|
|
if isinstance(t, Enum): |
|
|
return 'int' |
|
|
if isinstance(t, TypeDecl): |
|
|
return format_type(t.type) |
|
|
if isinstance(t, IdentifierType): |
|
|
assert len(t.names) == 1, f'Expected a single name, got {t.names}' |
|
|
return __c_type_to_python_type.get(t.names[0]) or 'ffi.CData' |
|
|
return t.name |
|
|
|
|
|
class PythonStubFuncDeclVisitor(c_ast.NodeVisitor): |
|
|
def __init__(self): |
|
|
self.sigs = {} |
|
|
self.sources = {} |
|
|
|
|
|
def get_source_snippet_lines(self, coord: pycparser.plyparser.Coord) -> Tuple[list[str], list[str]]: |
|
|
if coord.file not in self.sources: |
|
|
with open(coord.file, 'rt') as f: |
|
|
self.sources[coord.file] = f.readlines() |
|
|
source_lines = self.sources[coord.file] |
|
|
ncomment_lines = len(list(itertools.takewhile(lambda i: re.search(r'^\s*(//|/\*)', source_lines[i]), range(coord.line - 2, -1, -1)))) |
|
|
comment_lines = [l.strip() for l in source_lines[coord.line - 1 - ncomment_lines:coord.line - 1]] |
|
|
decl_lines = [] |
|
|
for line in source_lines[coord.line - 1:]: |
|
|
decl_lines.append(line.rstrip()) |
|
|
if (';' in line) or ('{' in line): break |
|
|
return (comment_lines, decl_lines) |
|
|
|
|
|
def visit_Enum(self, node: Enum): |
|
|
if node.values is not None: |
|
|
for e in node.values.enumerators: |
|
|
self.sigs[e.name] = f' @property\n def {e.name}(self) -> int: ...' |
|
|
|
|
|
def visit_Typedef(self, node: Typedef): |
|
|
pass |
|
|
|
|
|
def visit_FuncDecl(self, node: FuncDecl): |
|
|
ret_type = node.type |
|
|
is_ptr = False |
|
|
while isinstance(ret_type, PtrDecl): |
|
|
ret_type = ret_type.type |
|
|
is_ptr = True |
|
|
|
|
|
fun_name = ret_type.declname |
|
|
if fun_name.startswith('__'): |
|
|
return |
|
|
|
|
|
args = [] |
|
|
argnames = [] |
|
|
def gen_name(stem): |
|
|
i = 1 |
|
|
while True: |
|
|
new_name = stem if i == 1 else f'{stem}{i}' |
|
|
if new_name not in argnames: return new_name |
|
|
i += 1 |
|
|
|
|
|
for a in node.args.params: |
|
|
if isinstance(a, EllipsisParam): |
|
|
arg_name = gen_name('args') |
|
|
argnames.append(arg_name) |
|
|
args.append('*' + gen_name('args')) |
|
|
elif format_type(a.type) == 'None': |
|
|
continue |
|
|
else: |
|
|
arg_name = a.name or gen_name('arg') |
|
|
argnames.append(arg_name) |
|
|
args.append(f'{arg_name}: {format_type(a.type)}') |
|
|
|
|
|
ret = format_type(ret_type if not is_ptr else node.type) |
|
|
|
|
|
comment_lines, decl_lines = self.get_source_snippet_lines(node.coord) |
|
|
|
|
|
lines = [f' def {fun_name}({", ".join(args)}) -> {ret}:'] |
|
|
if len(comment_lines) == 0 and len(decl_lines) == 1: |
|
|
lines += [f' """{decl_lines[0]}"""'] |
|
|
else: |
|
|
lines += [' """'] |
|
|
lines += [f' {c.lstrip("/* ")}' for c in comment_lines] |
|
|
if len(comment_lines) > 0: |
|
|
lines += [''] |
|
|
lines += [f' {d}' for d in decl_lines] |
|
|
lines += [' """'] |
|
|
lines += [' ...'] |
|
|
self.sigs[fun_name] = '\n'.join(lines) |
|
|
|
|
|
def generate_stubs(header: str): |
|
|
""" |
|
|
Generates a .pyi Python stub file for the GGML API using C header files. |
|
|
""" |
|
|
|
|
|
v = PythonStubFuncDeclVisitor() |
|
|
v.visit(CParser().parse(header, "<input>")) |
|
|
|
|
|
keys = list(v.sigs.keys()) |
|
|
keys.sort() |
|
|
|
|
|
return '\n'.join([ |
|
|
'# auto-generated file', |
|
|
'import ggml.ffi as ffi', |
|
|
'import numpy as np', |
|
|
'class lib:', |
|
|
*[v.sigs[k] for k in keys] |
|
|
]) |
|
|
|
