| """ |
| 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] |
| ]) |
|
|
