chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Parse all headers for possible functions.

Browse Source 
Allows other files (e.g. mjrfilament) to define functions that
will be added to the introspection output.

PiperOrigin-RevId: 932403177
Change-Id: Ie2d2c6be67f0706463fdec4fc498364125e05dad
This commit is contained in:
Haroon Qureshi
2026-06-15 05:39:11 -07:00
committed by Copybara-Service
parent 81bc303a8f
commit b1f4656d1a
1 changed files with 16 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/introspect/codegen/generate_functions.py
+16 -7
View File
@@ -19,6 +19,7 @@ The JSON input can be generated via:
""" """
import json import json
import os
from typing import Any, Mapping, Sequence from typing import Any, Mapping, Sequence
from absl import app from absl import app
@@ -28,8 +29,8 @@ from introspect import ast_nodes
from introspect import type_parsing from introspect import type_parsing
from . import formatter from . import formatter
_HEADER_PATH = flags.DEFINE_string( _HEADER_PATHS = flags.DEFINE_string(
'header_path', None, 'Path to the original mujoco.h') 'header_paths', None, 'Path to all header files')
_JSON_PATH = flags.DEFINE_string( _JSON_PATH = flags.DEFINE_string(
'json_path', None, 'json_path', None,
'Path to the JSON file representing the Clang AST for mujoco.h') 'Path to the JSON file representing the Clang AST for mujoco.h')
@@ -47,8 +48,9 @@ def traverse(node, visitor):
class MjFunctionVisitor: class MjFunctionVisitor:
"""A Clang AST JSON node visitor for MuJoCo API function declarations.""" """A Clang AST JSON node visitor for MuJoCo API function declarations."""
def __init__(self, raw_header): def __init__(self, raw_headers: Mapping[str, str]):
self._raw_header = raw_header self._raw_headers = raw_headers
self._current_source = ''
self._functions = {} self._functions = {}
def _make_function(self, node: ClangJsonNode) -> ast_nodes.FunctionDecl: def _make_function(self, node: ClangJsonNode) -> ast_nodes.FunctionDecl:
@@ -101,7 +103,7 @@ class MjFunctionVisitor:
if type_name.endswith('*'): if type_name.endswith('*'):
decl_begin = node['range']['begin']['offset'] decl_begin = node['range']['begin']['offset']
decl_end = node['range']['end']['offset'] + node['range']['end']['tokLen'] decl_end = node['range']['end']['offset'] + node['range']['end']['tokLen']
decl = self._raw_header[decl_begin:decl_end] decl = self._raw_headers[self._current_source][decl_begin:decl_end]
name_begin = node['loc']['offset'] - decl_begin name_begin = node['loc']['offset'] - decl_begin
name_end = name_begin + node['loc']['tokLen'] name_end = name_begin + node['loc']['tokLen']
type_name = decl[:name_begin] + decl[name_end:] type_name = decl[:name_begin] + decl[name_end:]
@@ -128,6 +130,9 @@ class MjFunctionVisitor:
return ''.join(strings) return ''.join(strings)
def visit(self, node: ClangJsonNode) -> None: def visit(self, node: ClangJsonNode) -> None:
if 'loc' in node and 'file' in node['loc']:
self._current_source = os.path.basename(node['loc']['file'])
# Skip mjs_setUserValueWithCleanup as it's only useful for heap allocated # Skip mjs_setUserValueWithCleanup as it's only useful for heap allocated
# objects and doesn't need a python wrapper. # objects and doesn't need a python wrapper.
if ( if (
@@ -150,8 +155,12 @@ def main(argv: Sequence[str]) -> None:
with open(_JSON_PATH.value, 'r', encoding='utf-8') as f: with open(_JSON_PATH.value, 'r', encoding='utf-8') as f:
root = json.load(f) root = json.load(f)
with open(_HEADER_PATH.value, 'r') as f: raw_headers = {}
visitor = MjFunctionVisitor(f.read()) for p in (_HEADER_PATHS.value or '').split():
with open(p, 'r') as f:
raw_headers[os.path.basename(p)] = f.read()
visitor = MjFunctionVisitor(raw_headers)
traverse(root, visitor) traverse(root, visitor)