initial pass

generator/parser.py

@@ -0,0 +1,152 @@
+from dataclasses import dataclass
+from typing import Optional, TypedDict
+from pycparser import c_ast, parse_file, preprocess_file
+from pathlib import Path
+import os
+import json
+import shutil
+import logging
+
+logger = logging.getLogger(__name__)
+
+class ExtractedStructAttributeUnion(TypedDict):
+    type: Optional[str]
+
+class ExtractedStructAttribute(TypedDict):
+    type: Optional[str]
+    union: Optional[dict[str, Optional[str]]]
+
+ExtractedStruct = dict[str, ExtractedStructAttribute]
+ExtractedEnum = dict[str, Optional[str]]
+ExtractedFunctionParam = tuple[str, Optional[str]]
+
+class ExtractedFunction(TypedDict):
+    return_type: Optional[str]
+    params: list[ExtractedFunctionParam]
+
+@dataclass
+class ExtractedSymbols:
+    structs: dict[str, ExtractedStruct]
+    enums: dict[str, ExtractedEnum]
+    functions: dict[str, ExtractedFunction]
+
+def get_type_names(node: c_ast.Node, prefix: str="") -> Optional[str]:
+    if isinstance(node, c_ast.TypeDecl) and hasattr(node, 'quals') and node.quals:
+        prefix = " ".join(node.quals) + " " + prefix
+    if isinstance(node, c_ast.PtrDecl):
+        prefix = "*" + prefix
+
+    if hasattr(node, 'names'):
+        return prefix + node.names[0] # type: ignore
+    elif hasattr(node, 'type'):
+        return get_type_names(node.type, prefix) # type: ignore
+    return None
+
+class Visitor(c_ast.NodeVisitor):
+    def __init__(self):
+        self.structs: dict[str, ExtractedStruct] = {}
+        self.enums: dict[str, ExtractedEnum] = {}
+        self.functions: dict[str, ExtractedFunction] = {}
+
+    def visit_FuncDecl(self, node: c_ast.FuncDecl):
+        # node.show()
+        logger.debug(node)
+        node_type = node.type
+        is_pointer = False
+        if isinstance(node.type, c_ast.PtrDecl):
+            node_type = node.type.type
+            is_pointer = True
+        
+        if hasattr(node_type, "declname"):
+            return_type = get_type_names(node_type.type)
+            if return_type is not None and is_pointer:
+                return_type = "*" + return_type
+            func: ExtractedFunction = {
+                'return_type': return_type,
+                'params': [],
+            }
+            for param in node.args.params:
+                if param.name is None:
+                    continue
+                func['params'].append((param.name, get_type_names(param)))
+            self.functions[node_type.declname] = func
+        self.generic_visit(node)
+
+    def visit_Struct(self, node: c_ast.Struct):
+        # node.show()
+        if node.name and node.decls:
+            struct = {}
+            for decl in node.decls:
+                struct[decl.name] = {
+                    "type": get_type_names(decl),
+                }
+            self.structs[node.name] = struct
+        self.generic_visit(node)
+
+    def visit_Typedef(self, node: c_ast.Typedef):
+        # node.show()
+        if hasattr(node.type, 'type') and hasattr(node.type.type, 'decls') and node.type.type.decls:
+            struct = {}
+            for decl in node.type.type.decls:
+                if hasattr(decl, 'type') and hasattr(decl.type, 'type') and isinstance(decl.type.type, c_ast.Union):
+                    union = {}
+                    for field in decl.type.type.decls:
+                        union[field.name] = get_type_names(field)
+                    struct[decl.name] = {
+                        'union': union
+                    }
+                else:
+                    struct[decl.name] = {
+                        "type": get_type_names(decl),
+                    }
+            self.structs[node.name] = struct
+        if hasattr(node.type, 'type') and isinstance(node.type.type, c_ast.Enum):
+            enum = {}
+            for enumerator in node.type.type.values.enumerators:
+                if enumerator.value is None:
+                    enum[enumerator.name] = None
+                else:
+                    enum[enumerator.name] = enumerator.value.value
+            self.enums[node.name] = enum
+        self.generic_visit(node)
+
+
+def parse_headers(input_files: list[Path], output_dir: Path, tmp_dir: Path) -> ExtractedSymbols:
+    cpp_args = ["-nostdinc", "-D__attribute__(x)=", "-E"]
+
+    # Make a new clay.h that combines the provided input files, so that we can add bindings for customized structs
+    with open(tmp_dir / 'merged_clay.h', 'w') as f:
+        for input_file in input_files:
+            with open(input_file, 'r') as f2:
+                for line in f2:
+                    # Ignore includes, as they should be manually included in input_files.
+                    if line.startswith("#include"):
+                        continue
+
+                    # Ignore the CLAY_IMPLEMENTATION define, because we only want to parse the public api code.
+                    # This is helpful so that the user can provide their implementation code, which will contain any custom extensions
+                    if "#define CLAY_IMPLEMENTATION" in line:
+                        continue
+
+                    f.write(line)
+
+    # Preprocess the file
+    logger.info("Preprocessing file")
+    preprocessed = preprocess_file(tmp_dir / 'merged_clay.h', cpp_path="cpp", cpp_args=cpp_args) # type: ignore
+    with open(tmp_dir / 'clay.preprocessed.h', 'w') as f:
+        f.write(preprocessed)
+
+    # Parse the file
+    logger.info("Parsing file")
+    ast = parse_file(tmp_dir / 'clay.preprocessed.h', use_cpp=False) # type: ignore
+
+    # Extract symbols
+    visitor = Visitor()
+    visitor.visit(ast)
+    
+    result = ExtractedSymbols(
+        structs=visitor.structs,
+        enums=visitor.enums,
+        functions=visitor.functions
+    )
+    return result