Kyle Bayes
d3d789cf7d
PiperOrigin-RevId: 515431052 Change-Id: I440de0cb417ce216b6d636a5d86dc9ef6cc1fc06
114 lines
3.5 KiB
Python
114 lines
3.5 KiB
Python
# 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.
|
|
# ==============================================================================
|
|
"""Code generator for function_traits.h."""
|
|
|
|
from typing import Mapping, Sequence
|
|
|
|
from absl import app
|
|
|
|
from introspect import ast_nodes
|
|
from introspect import functions
|
|
|
|
FUNCTIONS: Mapping[str, ast_nodes.FunctionDecl] = functions.FUNCTIONS
|
|
|
|
|
|
def main(argv: Sequence[str]) -> None:
|
|
if len(argv) > 1:
|
|
raise app.UsageError('Too many command-line arguments.')
|
|
|
|
struct_decls = []
|
|
for func in FUNCTIONS.values():
|
|
# Skip mju_error_{i,s} and mju_warning_{i,s} as these are not
|
|
# supported in the Python bindings, and Introspect currently
|
|
# doesn't support variadic functions.
|
|
if func.name.startswith('mju_error') or func.name == 'mju_warning':
|
|
continue
|
|
|
|
# Modify some parameter types.
|
|
parameters = []
|
|
modified = False
|
|
for p in func.parameters:
|
|
# Expose array parameters as pointer-to-arrays so that we can determine
|
|
# array extents in C++ templates.
|
|
if isinstance(p.type, ast_nodes.ArrayType):
|
|
parameters.append(ast_nodes.FunctionParameterDecl(
|
|
name=p.name, type=ast_nodes.PointerType(
|
|
ast_nodes.ArrayType(
|
|
inner_type=p.type.inner_type, extents=p.type.extents))))
|
|
modified = True
|
|
else:
|
|
parameters.append(p)
|
|
|
|
if modified:
|
|
func = ast_nodes.FunctionDecl(
|
|
name=func.name, return_type=func.return_type,
|
|
parameters=parameters, doc=func.doc)
|
|
getfunc = f'*reinterpret_cast<type*>(&::{func.name})'
|
|
else:
|
|
getfunc = f'::{func.name}'
|
|
|
|
param_names = ', '.join(f'"{p.name}"' for p in parameters)
|
|
|
|
struct_decls.append(f"""
|
|
struct {func.name} {{
|
|
static constexpr char name[] = "{func.name}";
|
|
static constexpr char doc[] = "{func.doc}";
|
|
using type = {func.decltype};
|
|
static constexpr auto param_names = std::make_tuple({param_names});
|
|
|
|
MUJOCO_ALWAYS_INLINE static type& GetFunc() {{
|
|
return {getfunc};
|
|
}}
|
|
}};
|
|
""".strip())
|
|
|
|
all_structs = '\n\n'.join(struct_decls)
|
|
|
|
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.
|
|
|
|
#ifndef MUJOCO_PYTHON_CODEGEN_FUNCTION_TRAITS_H_
|
|
#define MUJOCO_PYTHON_CODEGEN_FUNCTION_TRAITS_H_
|
|
|
|
#include <tuple>
|
|
|
|
#include <mujoco/mujoco.h>
|
|
#include "util/crossplatform.h"
|
|
|
|
namespace mujoco::python_traits {{
|
|
|
|
{all_structs}
|
|
|
|
}} // namespace mujoco::python_traits
|
|
|
|
#endif // MUJOCO_PYTHON_CODEGEN_FUNCTION_TRAITS_H_
|
|
""".lstrip())
|
|
|
|
|
|
if __name__ == '__main__':
|
|
app.run(main)
|