Version 2.1.2: Python bindings, OBJ assets support, bugfixes.

PiperOrigin-RevId: 434731612
Change-Id: I0cfda3e7a3d1c72036764986efc252ffa1b8c6b0

python/mujoco/codegen/generate_function_traits.py

@@ -0,0 +1,108 @@
+# 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():
+
+    # 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.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)