seamless_communication/ggml/examples/python/stubs.py

"""
  This generates .pyi stubs for the cffi Python bindings generated by regenerate.py
"""
import sys, re, itertools
sys.path.extend(['.', '..']) # for pycparser

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