# Copyright 2017 The dm_control Authors.

#

# 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.

# ============================================================================

r"""Automatically generates ctypes Python bindings for MuJoCo.

Parses the following MuJoCo header files:

mjdata.h

mjmodel.h

mjrender.h

mjui.h

mjvisualize.h

mjxmacro.h

mujoco.h;

generates the following Python source files:

constants.py: constants

enums.py: enums

sizes.py: size information for dynamically-shaped arrays

Example usage:

autowrap --header_paths='/path/to/mjmodel.h /path/to/mjdata.h ...' \

--output_dir=/path/to/mjbindings

"""

import collections

import io

import os

from absl import app

from absl import flags

from absl import logging

from dm_control.autowrap import binding_generator

from dm_control.autowrap import codegen_util

_MJMODEL_H = "mjmodel.h"

_MJXMACRO_H = "mjxmacro.h"

FLAGS = flags.FLAGS

flags.DEFINE_spaceseplist(

"header_paths", None,

"Space-separated list of paths to MuJoCo header files.")

flags.DEFINE_string("output_dir", None,

"Path to output directory for wrapper source files.")

def main(unused_argv):

special_header_paths = {}

# Get the path to the mjmodel and mjxmacro header files.

# These header files need special handling.

for header in (_MJMODEL_H, _MJXMACRO_H):

for path in FLAGS.header_paths:

if path.endswith(header):

special_header_paths[header] = path

break

if header not in special_header_paths:

logging.fatal("List of inputs must contain a path to %s", header)

# Make sure mjmodel.h is parsed first, since it is included by other headers.

srcs = codegen_util.UniqueOrderedDict()

sorted_header_paths = sorted(FLAGS.header_paths)

sorted_header_paths.remove(special_header_paths[_MJMODEL_H])

sorted_header_paths.insert(0, special_header_paths[_MJMODEL_H])

for p in sorted_header_paths:

with io.open(p, "r", errors="ignore") as f:

srcs[p] = f.read()

# consts_dict should be a codegen_util.UniqueOrderedDict.

# This is a temporary workaround due to the fact that the parser does not yet

# handle nested `#if define(predicate)` blocks, which results in some

# constants being parsed twice. We therefore can't enforce the uniqueness of

# the keys in `consts_dict`. As of MuJoCo v1.30 there is only a single problem

# block beginning on line 10 in mujoco.h, and a single constant is affected

# (MJAPI).

consts_dict = collections.OrderedDict()

# These are commented in `mjdata.h` but have no macros in `mjxmacro.h`.

hints_dict = codegen_util.UniqueOrderedDict({"buffer": ("nbuffer",),

"stack": ("nstack",)})

parser = binding_generator.BindingGenerator(

consts_dict=consts_dict, hints_dict=hints_dict)

# Parse enums.

for pth, src in srcs.items():

if pth is not special_header_paths[_MJXMACRO_H]:

parser.parse_enums(src)

# Parse constants and type declarations.

for pth, src in srcs.items():

if pth is not special_header_paths[_MJXMACRO_H]:

parser.parse_consts_typedefs(src)

# Get shape hints from mjxmacro.h.

parser.parse_hints(srcs[special_header_paths[_MJXMACRO_H]])

# Create the output directory if it doesn't already exist.

if not os.path.exists(FLAGS.output_dir):

os.makedirs(FLAGS.output_dir)

# Generate Python source files and write them to the output directory.

parser.write_consts(os.path.join(FLAGS.output_dir, "constants.py"))

parser.write_enums(os.path.join(FLAGS.output_dir, "enums.py"))

parser.write_index_dict(os.path.join(FLAGS.output_dir, "sizes.py"))

if __name__ == "__main__":

flags.mark_flag_as_required("header_paths")

flags.mark_flag_as_required("output_dir")

app.run(main)