From 71721ed8e6bf012e8fa7c9eefc4d688a4858e94f Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Fri, 19 Jun 2026 01:16:35 -0700 Subject: [PATCH] Merge generate scripts into a single script. PiperOrigin-RevId: 934787488 Change-Id: I1e4d5870a7f4b339fa0de60271211fdffe80a3d8 --- python/mujoco/introspect/codegen/generate.py | 465 ++++++++++++++++++ .../introspect/codegen/generate_enums.py | 128 ----- .../introspect/codegen/generate_functions.py | 202 -------- .../introspect/codegen/generate_structs.py | 283 ----------- 4 files changed, 465 insertions(+), 613 deletions(-) create mode 100644 python/mujoco/introspect/codegen/generate.py delete mode 100644 python/mujoco/introspect/codegen/generate_enums.py delete mode 100644 python/mujoco/introspect/codegen/generate_functions.py delete mode 100644 python/mujoco/introspect/codegen/generate_structs.py diff --git a/python/mujoco/introspect/codegen/generate.py b/python/mujoco/introspect/codegen/generate.py new file mode 100644 index 00000000..78b2cb6b --- /dev/null +++ b/python/mujoco/introspect/codegen/generate.py @@ -0,0 +1,465 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Generates enums.py, functions.py, and structs.py. + +The JSON input can be generated via: + clang -Xclang -ast-dump=json -fsyntax-only -fparse-all-comments -x c mujoco.h +""" + +import itertools +import json +import os +import re +from typing import Any, Mapping, Sequence, Union + +from absl import app +from absl import flags + +from introspect import ast_nodes +from introspect import type_parsing +from . import formatter + + +_JSON_PATH = flags.DEFINE_string( + 'json_path', None, + 'Path to the JSON file representing the Clang AST for mujoco.h') +_HEADER_PATHS = flags.DEFINE_string( + 'header_paths', None, 'Path to all header files') +_OUT_STRUCTS = flags.DEFINE_string( + 'out_structs', None, 'Path to the output structs file') +_OUT_FUNCTIONS = flags.DEFINE_string( + 'out_functions', None, 'Path to the output functions file') +_OUT_ENUMS = flags.DEFINE_string( + 'out_enums', None, 'Path to the output enums file') + + +_ANONYMOUS_KEY_PATTERN = re.compile(r'\d+:\d+(?=\))') + +_EXCLUDED = ( + 'mjpDecoder', + 'mjpDecoder_', + 'mjpEncoder', + 'mjpEncoder_', + 'mjpPlugin', + 'mjpPlugin_', + 'mjpResourceProvider', + 'mjpResourceProvider_', + 'mjResource', + 'mjResource_', + # Skip mjs_setUserValueWithCleanup as it's only useful for heap allocated + # objects and doesn't need a python wrapper. + 'mjs_setUserValueWithCleanup', +) + +_ARRAY_COMMENT_PATTERN = re.compile(r'(.+?)\s+\(([^\(\)]+) x ([^\(\)]+)\)\Z') + +_HEADER_TEMPLATE = ''' +# Copyright {year} DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Provides information about MuJoCo API {type}. + +DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED. +""" +'''.strip() + +_ENUMS_TEMPLATE = ''' + +from typing import Mapping + +from .ast_nodes import EnumDecl + +ENUMS: Mapping[str, EnumDecl] = {enums_str} +''' + + +_FUNCTIONS_TEMPLATE = ''' + +from typing import Mapping + +from .ast_nodes import ArrayType +from .ast_nodes import FunctionDecl +from .ast_nodes import FunctionParameterDecl +from .ast_nodes import PointerType +from .ast_nodes import ValueType + +FUNCTIONS: Mapping[str, FunctionDecl] = {functions_str} +''' + +_STRUCTS_TEMPLATE = ''' + +from typing import Mapping + +from .ast_nodes import AnonymousStructDecl +from .ast_nodes import AnonymousUnionDecl +from .ast_nodes import ArrayType +from .ast_nodes import PointerType +from .ast_nodes import StructDecl +from .ast_nodes import StructFieldDecl +from .ast_nodes import ValueType + +STRUCTS: Mapping[str, StructDecl] = {structs_str} +''' + + +ClangJsonNode = Mapping[str, Any] + + +class _AnonymousTypePlaceholder(ast_nodes.ValueType): + + def __init__(self, anonymous_key: str): + self.name = anonymous_key + self.is_const = False + self.is_volatile = False + + +class NodeVisitor: + """A Clang AST JSON node visitor for.""" + + def __init__(self, raw_headers: Mapping[str, str]): + self._enums = {} + self._exported_enums = {} + self._raw_headers = raw_headers + self._current_source = '' + self._exported_functions = {} + self._structs = {} + self._anonymous = {} + self._exported_structs = {} + + @property + def exported_enums(self) -> Mapping[str, ast_nodes.EnumDecl]: + return self._exported_enums + + @property + def exported_functions(self) -> Mapping[str, ast_nodes.FunctionDecl]: + return self._exported_functions + + @property + def exported_structs(self) -> Mapping[str, ast_nodes.StructDecl]: + return self._exported_structs + + def visit(self, node: ClangJsonNode) -> None: + """Visits a JSON node and stores information about API elements.""" + if 'loc' in node and 'file' in node['loc']: + self._current_source = os.path.basename(node['loc']['file']) + + node_kind = node.get('kind') + node_name = node.get('name', '') + + if ( + node_kind == 'FunctionDecl' + and node_name.startswith('mj') + and node_name not in _EXCLUDED + ): + func_decl = self._make_function(node) + self._exported_functions[func_decl.name] = func_decl + elif node_kind == 'EnumDecl' and node_name.startswith('mj'): + enum_decl = self._make_enum(node) + self._enums[enum_decl.name] = enum_decl + elif node_kind == 'RecordDecl' and self._is_mujoco_type(node): + struct_decl = self._make_struct(node) + if hasattr(struct_decl, 'name'): + self._structs[struct_decl.name] = struct_decl + else: + anonymous_key = self._make_anonymous_key(node) + if anonymous_key in self._anonymous: + raise RuntimeError( + f'duplicate key for anonymous struct: {anonymous_key}') + self._anonymous[anonymous_key] = struct_decl + elif (node_kind == 'TypedefDecl' and + node['type']['qualType'].startswith('enum mj')): + enum = self._enums[node['type']['qualType']] + self._exported_enums[node_name] = ast_nodes.EnumDecl( + name=node_name, declname=enum.declname, values=dict(enum.values)) + elif (node_kind == 'TypedefDecl' and + node['type']['qualType'].startswith('struct mj') and + node_name not in _EXCLUDED): + declname = node['type']['qualType'] + try: + struct = self._structs[declname] + except KeyError: + self._exported_structs[node_name] = ast_nodes.StructDecl( + name=node_name, declname=declname, fields=()) + else: + self._exported_structs[node_name] = ast_nodes.StructDecl( + name=node_name, declname=struct.declname, fields=struct.fields) + + def resolve(self) -> None: + """Replaces anonymous struct placeholders with corresponding decl.""" + for struct in itertools.chain( + self._structs.values(), self._exported_structs.values() + ): + fields = [] + for field in struct.fields: + if isinstance(field, _AnonymousTypePlaceholder): + fields.append(self._anonymous[field.name]) + elif isinstance(field.type, _AnonymousTypePlaceholder): + fields.append( + ast_nodes.StructFieldDecl( + name=field.name, + type=self._anonymous[field.type.name], + doc=field.doc, + ) + ) + else: + fields.append(field) + struct.fields = tuple(fields) + + def _make_enum(self, node: ClangJsonNode) -> ast_nodes.EnumDecl: + """Makes a EnumDecl from a Clang AST EnumDecl node.""" + name = f"enum {node['name']}" + values = [] + for child in node['inner']: + child_kind = child.get('kind') + if child_kind == 'EnumConstantDecl': + next_idx = values[-1][1] + 1 if values else 0 + if 'inner' in child: + value = int(child['inner'][0].get('value', next_idx)) + else: + value = next_idx + values.append((child['name'], value)) + return ast_nodes.EnumDecl(name=name, declname=name, values=dict(values)) + + def _make_function(self, node: ClangJsonNode) -> ast_nodes.FunctionDecl: + """Makes a FunctionDecl from a Clang AST FunctionDecl node.""" + name = node['name'] + return_type = type_parsing.parse_function_return_type( + node['type']['qualType']) + parameters = [] + comments = [] + nullable_params = set() + + for child in node['inner']: + child_kind = child.get('kind') + if child_kind == 'FullComment': + comments.append(self._make_function_comment(child)) + nullable_params.update(self._find_nullable_params(child)) + comment = ' '.join(comments).strip() + + for child in node['inner']: + child_kind = child.get('kind') + if child_kind == 'ParmVarDecl': + parameters.append(self._make_parameter(child, nullable_params)) + + return ast_nodes.FunctionDecl( + name=name, return_type=return_type, parameters=parameters, doc=comment) + + def _find_nullable_params(self, node: ClangJsonNode) -> set[str]: + """Finds the names of parameters that are marked as nullable.""" + nullable_params = set() + for child in node['inner']: + child_kind = child.get('kind') + if child_kind == 'ParagraphComment': + nullable_params.update(self._find_nullable_params(child)) + if child_kind == 'TextComment': + if 'Nullable' in child['text']: + for param in child['text'].split(':')[1].split(','): + nullable_params.add(param.strip()) + return nullable_params + + def _make_parameter( + self, node: ClangJsonNode, nullable_params: set[str] + ) -> ast_nodes.FunctionParameterDecl: + """Makes a ParameterDecl from a Clang AST ParmVarDecl node.""" + name = node['name'] + type_name = node['type']['qualType'] + nullable = name in nullable_params + + # For a pointer parameters, look up in the original header to see if + # n array extent was declared there. + if type_name.endswith('*'): + decl_begin = node['range']['begin']['offset'] + decl_end = node['range']['end']['offset'] + node['range']['end']['tokLen'] + decl = self._raw_headers[self._current_source][decl_begin:decl_end] + name_begin = node['loc']['offset'] - decl_begin + name_end = name_begin + node['loc']['tokLen'] + type_name = decl[:name_begin] + decl[name_end:] + + return ast_nodes.FunctionParameterDecl( + nullable=nullable, + name=name, + type=type_parsing.parse_type(type_name), + ) + + def _make_function_comment(self, node: ClangJsonNode) -> str: + """Makes a comment string from a Clang AST FullComment node.""" + kind = node.get('kind') + if kind == 'TextComment': + return node['text'].replace('\N{NO-BREAK SPACE}', ' ') + else: + strings = [] + for child in node['inner']: + comment = self._make_function_comment(child) + nullable_index = comment.find('Nullable:') + if nullable_index != -1: + comment = comment[:nullable_index] + strings.append(comment) + return ''.join(strings) + + def _make_struct_comment( + self, node: ClangJsonNode, strip: bool = True + ) -> str: + """Makes a comment string from a Clang AST FullComment node.""" + kind = node.get('kind') + if kind == 'TextComment': + retval = node['text'].replace('\N{NO-BREAK SPACE}', ' ') + else: + strings = [] + for child in node['inner']: + strings.append(self._make_struct_comment(child, strip=False)) + retval = ''.join(strings) + if strip: + retval = retval.strip() + return retval + + def _normalize_type( + self, declname: str + ) -> Union[ + ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]: + """Resolves anonymous structs/unions and looks up existing typedefs.""" + # Check for anonymous struct/union. + if '(unnamed ' in declname: + m = _ANONYMOUS_KEY_PATTERN.search(declname) + if not m: + raise RuntimeError('cannot parse anonymous key from {m!r}') + return _AnonymousTypePlaceholder(m.group(0)) + + # Lookup typedef name and use it instead if one exists. + for k, v in self._exported_structs.items(): + if declname == v.declname: + return type_parsing.parse_type(k) + + # No valid normalization, just parse the declname. + return type_parsing.parse_type(declname) + + def _make_field( + self, node: ClangJsonNode + ) -> Union[ast_nodes.StructFieldDecl, _AnonymousTypePlaceholder]: + """Makes a StructFieldDecl object from a Clang AST FieldDecl node.""" + doc = '' + for child in node.get('inner', ()): + if child['kind'] == 'FullComment': + doc = self._make_struct_comment(child) + if 'name' in node: + field_type = self._normalize_type(node['type']['qualType']) + m = _ARRAY_COMMENT_PATTERN.match(doc) + if m is None: + array_extent = None + else: + doc = m.group(1) + array_extent_0 = m.group(2) + array_extent_1 = m.group(3) + try: + array_extent_1 = int(array_extent_1) + except ValueError: + pass + if array_extent_1 == 1: + array_extent = (array_extent_0,) + else: + array_extent = (array_extent_0, array_extent_1) + return ast_nodes.StructFieldDecl( + name=node['name'], type=field_type, doc=doc, + array_extent=array_extent) + else: + return _AnonymousTypePlaceholder(self._make_anonymous_key(node)) + + def _make_struct( + self, node: ClangJsonNode + ) -> Union[ast_nodes.AnonymousStructDecl, ast_nodes.StructDecl]: + """Makes a Decl object from a Clang AST RecordDecl node.""" + name = f"{node['tagUsed']} {node['name']}" if 'name' in node else '' + fields = [] + for child in node.get('inner', ()): + child_kind = child.get('kind') + if child_kind == 'FieldDecl': + fields.append(self._make_field(child)) + + if name: + return ast_nodes.StructDecl(name=name, declname=name, fields=fields) + elif node['tagUsed'] == 'union': + return ast_nodes.AnonymousUnionDecl(fields=fields) + else: + return ast_nodes.AnonymousStructDecl(fields=fields) + + def _is_mujoco_type(self, node: ClangJsonNode) -> bool: + node_name = node.get('name', '') + included_from = os.path.basename( + node['loc'].get('includedFrom', {}).get('file', '') + ) + return node_name not in _EXCLUDED and ( + node_name.startswith('mj') + or included_from == 'mujoco.h' + or included_from.startswith('mj') + ) + + def _make_anonymous_key(self, node: ClangJsonNode) -> str: + line = node['loc']['line'] + col = node['loc']['col'] + return f'{line}:{col}' + + +def _traverse(node, visitor): + visitor.visit(node) + + children = node.get('inner', []) + for child in children: + _traverse(child, visitor) + + +def main(argv: Sequence[str]) -> None: + if len(argv) > 1: + raise app.UsageError('Too many command-line arguments.') + + with open(_JSON_PATH.value, 'r', encoding='utf-8') as f: + root = json.load(f) + + raw_headers = {} + for p in (_HEADER_PATHS.value or '').split(): + with open(p, 'r') as f: + raw_headers[os.path.basename(p)] = f.read() + + visitor = NodeVisitor(raw_headers) + _traverse(root, visitor) + visitor.resolve() + + with open(_OUT_FUNCTIONS.value, 'w') as f: + functions_str = formatter.format_as_python_code(visitor.exported_functions) + f.write(_HEADER_TEMPLATE.format(year=2022, type='functions')) + f.write(_FUNCTIONS_TEMPLATE.format(functions_str=functions_str)) + + with open(_OUT_ENUMS.value, 'w') as f: + enums_str = formatter.format_as_python_code(visitor.exported_enums) + f.write(_HEADER_TEMPLATE.format(year=2022, type='enums')) + f.write(_ENUMS_TEMPLATE.format(enums_str=enums_str)) + + with open(_OUT_STRUCTS.value, 'w') as f: + structs_str = formatter.format_as_python_code(visitor.exported_structs) + f.write(_HEADER_TEMPLATE.format(year=2023, type='structs')) + f.write(_STRUCTS_TEMPLATE.format(structs_str=structs_str)) + + +if __name__ == '__main__': + app.run(main) diff --git a/python/mujoco/introspect/codegen/generate_enums.py b/python/mujoco/introspect/codegen/generate_enums.py deleted file mode 100644 index 59d7ff3b..00000000 --- a/python/mujoco/introspect/codegen/generate_enums.py +++ /dev/null @@ -1,128 +0,0 @@ -# Copyright 2022 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Generates enums.py. - -The JSON input can be generated via: - clang -Xclang -ast-dump=json -fsyntax-only -fparse-all-comments -x c mujoco.h -""" - -import json -from typing import Any, Mapping, Sequence - -from absl import app -from absl import flags - -from introspect import ast_nodes -from . import formatter - -_JSON_PATH = flags.DEFINE_string( - 'json_path', None, - 'Path to the JSON file representing the Clang AST for mujoco.h') - -ClangJsonNode = Mapping[str, Any] - - -def traverse(node, visitor): - visitor.visit(node) - children = node.get('inner', []) - for child in children: - traverse(child, visitor) - - -class MjEnumVisitor: - """A Clang AST JSON node visitor for MuJoCo API enum declarations.""" - - def __init__(self): - self._enums = {} - self._typedefs = {} - - def _make_enum(self, node: ClangJsonNode) -> ast_nodes.EnumDecl: - """Makes a EnumDecl from a Clang AST EnumDecl node.""" - name = f"enum {node['name']}" - values = [] - for child in node['inner']: - child_kind = child.get('kind') - if child_kind == 'EnumConstantDecl': - next_idx = values[-1][1] + 1 if values else 0 - if 'inner' in child: - value = int(child['inner'][0].get('value', next_idx)) - else: - value = next_idx - values.append((child['name'], value)) - return ast_nodes.EnumDecl(name=name, declname=name, values=dict(values)) - - def visit(self, node: ClangJsonNode) -> None: - if (node.get('kind') == 'EnumDecl' and - node.get('name', '').startswith('mj')): - enum_decl = self._make_enum(node) - self._enums[enum_decl.name] = enum_decl - elif (node.get('kind') == 'TypedefDecl' and - node['type']['qualType'].startswith('enum mj')): - enum = self._enums[node['type']['qualType']] - self._typedefs[node['name']] = ast_nodes.EnumDecl( - name=node['name'], declname=enum.declname, values=dict(enum.values)) - - @property - def enums(self) -> Mapping[str, ast_nodes.EnumDecl]: - return self._enums - - @property - def typedefs(self) -> Mapping[str, ast_nodes.EnumDecl]: - return self._typedefs - - -def main(argv: Sequence[str]) -> None: - if len(argv) > 1: - raise app.UsageError('Too many command-line arguments.') - - with open(_JSON_PATH.value, 'r', encoding='utf-8') as f: - root = json.load(f) - - visitor = MjEnumVisitor() - - traverse(root, visitor) - - enums_str = formatter.format_as_python_code(visitor.typedefs) - - print(f''' -# Copyright 2022 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Provides information about MuJoCo API enums. - -DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED. -""" - -from typing import Mapping - -from .ast_nodes import EnumDecl - -ENUMS: Mapping[str, EnumDecl] = {enums_str} -'''.strip()) # `print` adds a trailing newline - - -if __name__ == '__main__': - app.run(main) diff --git a/python/mujoco/introspect/codegen/generate_functions.py b/python/mujoco/introspect/codegen/generate_functions.py deleted file mode 100644 index 8016a383..00000000 --- a/python/mujoco/introspect/codegen/generate_functions.py +++ /dev/null @@ -1,202 +0,0 @@ -# Copyright 2022 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Generates functions.py. - -The JSON input can be generated via: - clang -Xclang -ast-dump=json -fsyntax-only -fparse-all-comments -x c mujoco.h -""" - -import json -import os -from typing import Any, Mapping, Sequence - -from absl import app -from absl import flags - -from introspect import ast_nodes -from introspect import type_parsing -from . import formatter - -_HEADER_PATHS = flags.DEFINE_string( - 'header_paths', None, 'Path to all header files') -_JSON_PATH = flags.DEFINE_string( - 'json_path', None, - 'Path to the JSON file representing the Clang AST for mujoco.h') - -ClangJsonNode = Mapping[str, Any] - - -def traverse(node, visitor): - visitor.visit(node) - children = node.get('inner', []) - for child in children: - traverse(child, visitor) - - -class MjFunctionVisitor: - """A Clang AST JSON node visitor for MuJoCo API function declarations.""" - - def __init__(self, raw_headers: Mapping[str, str]): - self._raw_headers = raw_headers - self._current_source = '' - self._functions = {} - - def _make_function(self, node: ClangJsonNode) -> ast_nodes.FunctionDecl: - """Makes a FunctionDecl from a Clang AST FunctionDecl node.""" - name = node['name'] - return_type = type_parsing.parse_function_return_type( - node['type']['qualType']) - parameters = [] - comments = [] - nullable_params = set() - - for child in node['inner']: - child_kind = child.get('kind') - if child_kind == 'FullComment': - comments.append(self._make_comment(child)) - nullable_params.update(self._find_nullable_params(child)) - comment = ' '.join(comments).strip() - - for child in node['inner']: - child_kind = child.get('kind') - if child_kind == 'ParmVarDecl': - parameters.append(self._make_parameter(child, nullable_params)) - - return ast_nodes.FunctionDecl( - name=name, return_type=return_type, parameters=parameters, doc=comment) - - def _find_nullable_params(self, node: ClangJsonNode) -> set[str]: - """Finds the names of parameters that are marked as nullable.""" - nullable_params = set() - for child in node['inner']: - child_kind = child.get('kind') - if child_kind == 'ParagraphComment': - nullable_params.update(self._find_nullable_params(child)) - if child_kind == 'TextComment': - if 'Nullable' in child['text']: - for param in child['text'].split(':')[1].split(','): - nullable_params.add(param.strip()) - return nullable_params - - def _make_parameter( - self, node: ClangJsonNode, nullable_params: set[str] - ) -> ast_nodes.FunctionParameterDecl: - """Makes a ParameterDecl from a Clang AST ParmVarDecl node.""" - name = node['name'] - type_name = node['type']['qualType'] - nullable = name in nullable_params - - # For a pointer parameters, look up in the original header to see if - # n array extent was declared there. - if type_name.endswith('*'): - decl_begin = node['range']['begin']['offset'] - decl_end = node['range']['end']['offset'] + node['range']['end']['tokLen'] - decl = self._raw_headers[self._current_source][decl_begin:decl_end] - name_begin = node['loc']['offset'] - decl_begin - name_end = name_begin + node['loc']['tokLen'] - type_name = decl[:name_begin] + decl[name_end:] - - return ast_nodes.FunctionParameterDecl( - nullable=nullable, - name=name, - type=type_parsing.parse_type(type_name), - ) - - def _make_comment(self, node: ClangJsonNode) -> str: - """Makes a comment string from a Clang AST FullComment node.""" - kind = node.get('kind') - if kind == 'TextComment': - return node['text'].replace('\N{NO-BREAK SPACE}', ' ') - else: - strings = [] - for child in node['inner']: - comment = self._make_comment(child) - nullable_index = comment.find('Nullable:') - if nullable_index != -1: - comment = comment[:nullable_index] - strings.append(comment) - return ''.join(strings) - - 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 - # objects and doesn't need a python wrapper. - if ( - node.get('kind') == 'FunctionDecl' - and node.get('name', '').startswith('mj') - and node.get('name', '') != 'mjs_setUserValueWithCleanup' - ): - func_decl = self._make_function(node) - self._functions[func_decl.name] = func_decl - - @property - def functions(self) -> Mapping[str, ast_nodes.FunctionDecl]: - return self._functions - - -def main(argv: Sequence[str]) -> None: - if len(argv) > 1: - raise app.UsageError('Too many command-line arguments.') - - with open(_JSON_PATH.value, 'r', encoding='utf-8') as f: - root = json.load(f) - - raw_headers = {} - 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) - - functions_str = formatter.format_as_python_code(visitor.functions) - - print(f''' -# Copyright 2022 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Provides information about MuJoCo API functions. - -DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED. -""" - -from typing import Mapping - -from .ast_nodes import ArrayType -from .ast_nodes import FunctionDecl -from .ast_nodes import FunctionParameterDecl -from .ast_nodes import PointerType -from .ast_nodes import ValueType - -FUNCTIONS: Mapping[str, FunctionDecl] = {functions_str} -'''.strip()) # `print` adds a trailing newline - - -if __name__ == '__main__': - app.run(main) diff --git a/python/mujoco/introspect/codegen/generate_structs.py b/python/mujoco/introspect/codegen/generate_structs.py deleted file mode 100644 index 6a060f65..00000000 --- a/python/mujoco/introspect/codegen/generate_structs.py +++ /dev/null @@ -1,283 +0,0 @@ -# Copyright 2023 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Generates structs.py. - -The JSON input can be generated via: - clang -Xclang -ast-dump=json -fsyntax-only -fparse-all-comments -x c mujoco.h -""" - -import itertools -import json -import os -import re -from typing import Any, Mapping, Sequence, Union - -from absl import app -from absl import flags - -from introspect import ast_nodes -from introspect import type_parsing -from . import formatter - -_JSON_PATH = flags.DEFINE_string( - 'json_path', None, - 'Path to the JSON file representing the Clang AST for mujoco.h') - -ClangJsonNode = Mapping[str, Any] - -_ANONYMOUS_KEY_PATTERN = re.compile(r'\d+:\d+(?=\))') -_EXCLUDED = ( - 'mjpDecoder', - 'mjpDecoder_', - 'mjpEncoder', - 'mjpEncoder_', - 'mjpPlugin', - 'mjpPlugin_', - 'mjpResourceProvider', - 'mjpResourceProvider_', - 'mjResource', - 'mjResource_', -) - -_ARRAY_COMMENT_PATTERN = re.compile(r'(.+?)\s+\(([^\(\)]+) x ([^\(\)]+)\)\Z') - - -def traverse(node, visitor): - visitor.visit(node) - children = node.get('inner', []) - for child in children: - traverse(child, visitor) - - -class _AnonymousTypePlaceholder(ast_nodes.ValueType): - - def __init__(self, anonymous_key: str): - self.name = anonymous_key - self.is_const = False - self.is_volatile = False - - -class MjStructVisitor: - """A Clang AST JSON node visitor for MuJoCo API struct declarations.""" - - def __init__(self): - self._structs = {} - self._anonymous = {} - self._typedefs = {} - - def _normalize_type( - self, declname: str - ) -> Union[ - ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]: - """Resolves anonymous structs/unions and looks up existing typedefs.""" - # Check for anonymous struct/union. - if '(unnamed ' in declname: - m = _ANONYMOUS_KEY_PATTERN.search(declname) - if not m: - raise RuntimeError('cannot parse anonymous key from {m!r}') - return _AnonymousTypePlaceholder(m.group(0)) - - # Lookup typedef name and use it instead if one exists. - for k, v in self._typedefs.items(): - if declname == v.declname: - return type_parsing.parse_type(k) - - # No valid normalization, just parse the declname. - return type_parsing.parse_type(declname) - - def _make_comment(self, node: ClangJsonNode, strip: bool = True) -> str: - """Makes a comment string from a Clang AST FullComment node.""" - kind = node.get('kind') - if kind == 'TextComment': - retval = node['text'].replace('\N{NO-BREAK SPACE}', ' ') - else: - strings = [] - for child in node['inner']: - strings.append(self._make_comment(child, strip=False)) - retval = ''.join(strings) - if strip: - retval = retval.strip() - return retval - - def _make_field( - self, node: ClangJsonNode - ) -> Union[ast_nodes.StructFieldDecl, _AnonymousTypePlaceholder]: - """Makes a StructFieldDecl object from a Clang AST FieldDecl node.""" - doc = '' - for child in node.get('inner', ()): - if child['kind'] == 'FullComment': - doc = self._make_comment(child) - if 'name' in node: - field_type = self._normalize_type(node['type']['qualType']) - m = _ARRAY_COMMENT_PATTERN.match(doc) - if m is None: - array_extent = None - else: - doc = m.group(1) - array_extent_0 = m.group(2) - array_extent_1 = m.group(3) - try: - array_extent_1 = int(array_extent_1) - except ValueError: - pass - if array_extent_1 == 1: - array_extent = (array_extent_0,) - else: - array_extent = (array_extent_0, array_extent_1) - return ast_nodes.StructFieldDecl( - name=node['name'], type=field_type, doc=doc, - array_extent=array_extent) - else: - return _AnonymousTypePlaceholder(self._make_anonymous_key(node)) - - def _make_struct( - self, node: ClangJsonNode - ) -> Union[ast_nodes.AnonymousStructDecl, ast_nodes.StructDecl]: - """Makes a Decl object from a Clang AST RecordDecl node.""" - name = f"{node['tagUsed']} {node['name']}" if 'name' in node else '' - fields = [] - for child in node.get('inner', ()): - child_kind = child.get('kind') - if child_kind == 'FieldDecl': - fields.append(self._make_field(child)) - - if name: - return ast_nodes.StructDecl(name=name, declname=name, fields=fields) - elif node['tagUsed'] == 'union': - return ast_nodes.AnonymousUnionDecl(fields=fields) - else: - return ast_nodes.AnonymousStructDecl(fields=fields) - - def _is_mujoco_type(self, node: ClangJsonNode) -> bool: - node_name = node.get('name', '') - included_from = os.path.basename( - node['loc'].get('includedFrom', {}).get('file', '') - ) - return node_name not in _EXCLUDED and ( - node_name.startswith('mj') - or included_from == 'mujoco.h' - or included_from.startswith('mj') - ) - - def _make_anonymous_key(self, node: ClangJsonNode) -> str: - line = node['loc']['line'] - col = node['loc']['col'] - return f'{line}:{col}' - - def visit(self, node: ClangJsonNode) -> None: - """Visits a JSON node.""" - if node.get('kind') == 'RecordDecl' and self._is_mujoco_type(node): - struct_decl = self._make_struct(node) - if hasattr(struct_decl, 'name'): - self._structs[struct_decl.name] = struct_decl - else: - anonymous_key = self._make_anonymous_key(node) - if anonymous_key in self._anonymous: - raise RuntimeError( - f'duplicate key for anonymous struct: {anonymous_key}') - self._anonymous[anonymous_key] = struct_decl - elif (node.get('kind') == 'TypedefDecl' and - node['type']['qualType'].startswith('struct mj') and - node['name'] not in _EXCLUDED): - declname = node['type']['qualType'] - try: - struct = self._structs[declname] - except KeyError: - self._typedefs[node['name']] = ast_nodes.StructDecl( - name=node['name'], declname=declname, fields=()) - else: - self._typedefs[node['name']] = ast_nodes.StructDecl( - name=node['name'], declname=struct.declname, fields=struct.fields) - - def resolve_all_anonymous(self) -> None: - """Replaces anonymous struct placeholders with corresponding decl.""" - for struct in itertools.chain( - self._structs.values(), self._typedefs.values() - ): - fields = [] - for field in struct.fields: - if isinstance(field, _AnonymousTypePlaceholder): - fields.append(self._anonymous[field.name]) - elif isinstance(field.type, _AnonymousTypePlaceholder): - fields.append( - ast_nodes.StructFieldDecl( - name=field.name, - type=self._anonymous[field.type.name], - doc=field.doc, - ) - ) - else: - fields.append(field) - struct.fields = tuple(fields) - - @property - def structs(self) -> Mapping[str, ast_nodes.StructDecl]: - return self._structs - - @property - def typedefs(self) -> Mapping[str, ast_nodes.StructDecl]: - return self._typedefs - - -def main(argv: Sequence[str]) -> None: - if len(argv) > 1: - raise app.UsageError('Too many command-line arguments.') - - with open(_JSON_PATH.value, 'r', encoding='utf-8') as f: - root = json.load(f) - - visitor = MjStructVisitor() - - traverse(root, visitor) - visitor.resolve_all_anonymous() - - structs_str = formatter.format_as_python_code(visitor.typedefs) - - print(f''' -# Copyright 2023 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Provides information about MuJoCo API structs. - -DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED. -""" - -from typing import Mapping - -from .ast_nodes import AnonymousStructDecl -from .ast_nodes import AnonymousUnionDecl -from .ast_nodes import ArrayType -from .ast_nodes import PointerType -from .ast_nodes import StructDecl -from .ast_nodes import StructFieldDecl -from .ast_nodes import ValueType - -STRUCTS: Mapping[str, StructDecl] = {structs_str} -'''.strip()) # `print` adds a trailing newline - - -if __name__ == '__main__': - app.run(main)