mujoco

# 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 re
from typing import Sequence

from absl import app
from absl import flags

from . import ast_processor
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}
'''


def main(argv: Sequence[str]) -> None:
  if len(argv) > 1:
    raise app.UsageError('Too many command-line arguments.')

  visitor = ast_processor.process(
      _JSON_PATH.value,  # pyrefly: ignore[bad-argument-type]
      (_HEADER_PATHS.value or '').split(),
      _EXCLUDED,
  )

  with open(_OUT_FUNCTIONS.value, 'w') as f:  # pyrefly: ignore[no-matching-overload]
    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:  # pyrefly: ignore[no-matching-overload]
    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:  # pyrefly: ignore[no-matching-overload]
    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)