Add mjSpec bindings.

Co-authored-by: Saran Tunyasuvunakool <stunya@google.com>
PiperOrigin-RevId: 648444641
Change-Id: I08ee1d3f4fae1efbc7934ba94702b1c06d8cb92a

doc/APIreference/functions.rst

@@ -4167,6 +4167,69 @@ mjs_nextChild
 
 Return body's next child of the same type; return NULL if child is last.
 
+.. _mjs_asBody:
+
+mjs_asBody
+~~~~~~~~~~
+
+.. mujoco-include:: mjs_asBody
+
+Safely cast an element as mjsBody, or return NULL if the element is not an mjsBody.
+
+.. _mjs_asGeom:
+
+mjs_asGeom
+~~~~~~~~~~
+
+.. mujoco-include:: mjs_asGeom
+
+Safely cast an element as mjsGeom, or return NULL if the element is not an mjsGeom.
+
+.. _mjs_asJoint:
+
+mjs_asJoint
+~~~~~~~~~~~
+
+.. mujoco-include:: mjs_asJoint
+
+Safely cast an element as mjsJoint, or return NULL if the element is not an mjsJoint.
+
+.. _mjs_asSite:
+
+mjs_asSite
+~~~~~~~~~~
+
+.. mujoco-include:: mjs_asSite
+
+Safely cast an element as mjsSite, or return NULL if the element is not an mjsSite.
+
+.. _mjs_asCamera:
+
+mjs_asCamera
+~~~~~~~~~~~~
+
+.. mujoco-include:: mjs_asCamera
+
+Safely cast an element as mjsCamera, or return NULL if the element is not an mjsCamera.
+
+.. _mjs_asLight:
+
+mjs_asLight
+~~~~~~~~~~~
+
+.. mujoco-include:: mjs_asLight
+
+Safely cast an element as mjsLight, or return NULL if the element is not an mjsLight.
+
+.. _mjs_asFrame:
+
+mjs_asFrame
+~~~~~~~~~~~
+
+.. mujoco-include:: mjs_asFrame
+
+Safely cast an element as mjsFrame, or return NULL if the element is not an mjsFrame.
+
 .. _AttributeSetters:
 
 Attribute setters

doc/changelog.rst

@@ -26,7 +26,6 @@ General
    Still missing:
 
    - Detailed documentation.
-   - Python bindings.
 
 .. youtube:: ZXBTEIDWHhs
    :align: right

doc/includes/references.h

@@ -3544,6 +3544,13 @@ mjsDefault* mjs_getSpecDefault(mjSpec* s);
 int mjs_getId(mjsElement* element);
 mjsElement* mjs_firstChild(mjsBody* body, mjtObj type);
 mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child);
+mjsBody* mjs_asBody(mjsElement* element);
+mjsGeom* mjs_asGeom(mjsElement* element);
+mjsJoint* mjs_asJoint(mjsElement* element);
+mjsSite* mjs_asSite(mjsElement* element);
+mjsCamera* mjs_asCamera(mjsElement* element);
+mjsLight* mjs_asLight(mjsElement* element);
+mjsFrame* mjs_asFrame(mjsElement* element);
 void mjs_setString(mjString* dest, const char* text);
 void mjs_setStringVec(mjStringVec* dest, const char* text);
 mjtByte mjs_setInStringVec(mjStringVec* dest, int i, const char* text);

include/mujoco/mujoco.h

@@ -1548,6 +1548,27 @@ MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type);
 // Return body's next child of the same type; return NULL if child is last.
 MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child);
 
+// Safely cast an element as mjsBody, or return NULL if the element is not an mjsBody.
+MJAPI mjsBody* mjs_asBody(mjsElement* element);
+
+// Safely cast an element as mjsGeom, or return NULL if the element is not an mjsGeom.
+MJAPI mjsGeom* mjs_asGeom(mjsElement* element);
+
+// Safely cast an element as mjsJoint, or return NULL if the element is not an mjsJoint.
+MJAPI mjsJoint* mjs_asJoint(mjsElement* element);
+
+// Safely cast an element as mjsSite, or return NULL if the element is not an mjsSite.
+MJAPI mjsSite* mjs_asSite(mjsElement* element);
+
+// Safely cast an element as mjsCamera, or return NULL if the element is not an mjsCamera.
+MJAPI mjsCamera* mjs_asCamera(mjsElement* element);
+
+// Safely cast an element as mjsLight, or return NULL if the element is not an mjsLight.
+MJAPI mjsLight* mjs_asLight(mjsElement* element);
+
+// Safely cast an element as mjsFrame, or return NULL if the element is not an mjsFrame.
+MJAPI mjsFrame* mjs_asFrame(mjsElement* element);
+
 
 //---------------------------------- Attribute setters ---------------------------------------------
 

introspect/codegen/generate_structs.py

@@ -166,9 +166,15 @@ def visit(self, node: ClangJsonNode) -> None:
     elif (node.get('kind') == 'TypedefDecl' and
           node['type']['qualType'].startswith('struct mj') and
           node['name'] not in _EXCLUDED):
-      struct = self._structs[node['type']['qualType']]
-      self._typedefs[node['name']] = ast_nodes.StructDecl(
-          name=node['name'], declname=struct.declname, fields=struct.fields)
+      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."""

introspect/functions.py

@@ -9813,6 +9813,118 @@
          ),
          doc="Return body's next child of the same type; return NULL if child is last.",  # pylint: disable=line-too-long
      )),
+    ('mjs_asBody',
+     FunctionDecl(
+         name='mjs_asBody',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsBody'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+         ),
+         doc='Safely cast an element as mjsBody, or return NULL if the element is not an mjsBody.',  # pylint: disable=line-too-long
+     )),
+    ('mjs_asGeom',
+     FunctionDecl(
+         name='mjs_asGeom',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsGeom'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+         ),
+         doc='Safely cast an element as mjsGeom, or return NULL if the element is not an mjsGeom.',  # pylint: disable=line-too-long
+     )),
+    ('mjs_asJoint',
+     FunctionDecl(
+         name='mjs_asJoint',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsJoint'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+         ),
+         doc='Safely cast an element as mjsJoint, or return NULL if the element is not an mjsJoint.',  # pylint: disable=line-too-long
+     )),
+    ('mjs_asSite',
+     FunctionDecl(
+         name='mjs_asSite',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsSite'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+         ),
+         doc='Safely cast an element as mjsSite, or return NULL if the element is not an mjsSite.',  # pylint: disable=line-too-long
+     )),
+    ('mjs_asCamera',
+     FunctionDecl(
+         name='mjs_asCamera',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsCamera'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+         ),
+         doc='Safely cast an element as mjsCamera, or return NULL if the element is not an mjsCamera.',  # pylint: disable=line-too-long
+     )),
+    ('mjs_asLight',
+     FunctionDecl(
+         name='mjs_asLight',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsLight'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+         ),
+         doc='Safely cast an element as mjsLight, or return NULL if the element is not an mjsLight.',  # pylint: disable=line-too-long
+     )),
+    ('mjs_asFrame',
+     FunctionDecl(
+         name='mjs_asFrame',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsFrame'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+         ),
+         doc='Safely cast an element as mjsFrame, or return NULL if the element is not an mjsFrame.',  # pylint: disable=line-too-long
+     )),
     ('mjs_setString',
      FunctionDecl(
          name='mjs_setString',

python/MANIFEST.in

@@ -1,3 +1,3 @@
 include LICENSE *.md
-recursive-include mujoco *.h *.cc *.mm CMakeLists.txt *.cmake
+recursive-include mujoco *.h *.cc *.cc.inc *.mm CMakeLists.txt *.cmake
 recursive-include mujoco/mjpython mjpython.* Info.plist

python/make_sdist.sh

@@ -44,6 +44,8 @@ python "${package_dir}"/mujoco/codegen/generate_enum_traits.py > \
     mujoco/enum_traits.h
 python "${package_dir}"/mujoco/codegen/generate_function_traits.py > \
     mujoco/function_traits.h
+python "${package_dir}"/mujoco/codegen/generate_spec_bindings.py > \
+    mujoco/specs.cc.inc
 export PYTHONPATH="${old_pythonpath}"
 
 # Copy over the LICENSE file.

python/mujoco/CMakeLists.txt

@@ -405,6 +405,29 @@ target_link_libraries(
           structs_header
 )
 
+if(NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/specs.cc.inc)
+  add_custom_command(
+    OUTPUT specs.cc.inc
+    COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH=${mujoco_SOURCE_DIR}/mujoco ${Python3_EXECUTABLE}
+            ${CMAKE_CURRENT_SOURCE_DIR}/codegen/generate_spec_bindings.py > specs.cc.inc
+    DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/codegen/generate_spec_bindings.py
+  )
+endif()
+
+mujoco_pybind11_module(
+  _specs
+  specs.cc
+  specs.cc.inc
+)
+target_link_libraries(
+  _specs
+  PRIVATE mujoco
+          Eigen3::Eigen
+          errors_header
+          raw
+          structs_header
+)
+
 mujoco_pybind11_module(_simulate simulate.cc)
 target_link_libraries(
   _simulate
@@ -424,6 +447,7 @@ set(LIBRARIES_FOR_WHEEL
     "$<TARGET_FILE:_render>"
     "$<TARGET_FILE:_rollout>"
     "$<TARGET_FILE:_simulate>"
+    "$<TARGET_FILE:_specs>"
     "$<TARGET_FILE:_structs>"
     "$<TARGET_FILE:mujoco>"
 )
@@ -459,6 +483,7 @@ if(MUJOCO_PYTHON_MAKE_WHEEL)
     _render
     _rollout
     _simulate
+    _specs
     _structs
     mujoco
   )

python/mujoco/__init__.py

@@ -43,6 +43,7 @@
 from mujoco._errors import *
 from mujoco._functions import *
 from mujoco._render import *
+from mujoco._specs import *
 from mujoco._structs import *
 from mujoco.gl_context import *
 from mujoco.renderer import Renderer