Add non-tree element lists to python bindings.

Fixes #1752.

PiperOrigin-RevId: 649405039
Change-Id: Ie0b575a3ec912b773141694397ac2f394bff1227

src/user/user_api.cc

@@ -590,6 +590,22 @@ mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child) {
 
 
 
+// return spec's first element of selected type
+mjsElement* mjs_firstElement(mjSpec* s, mjtObj type) {
+  mjCModel* modelC = static_cast<mjCModel*>(s->element);
+  return modelC->NextObject(NULL, type);
+}
+
+
+
+// return spec's next element; return NULL if element is last
+mjsElement* mjs_nextElement(mjSpec* s, mjsElement* element) {
+  mjCModel* modelC = static_cast<mjCModel*>(s->element);
+  return modelC->NextObject(element);
+}
+
+
+
 // return body given mjsElement
 mjsBody* mjs_asBody(mjsElement* element) {
   return element->elemtype == mjOBJ_BODY ? &(static_cast<mjCBody*>(element)->spec) : nullptr;
@@ -639,6 +655,118 @@ mjsFrame* mjs_asFrame(mjsElement* element) {
 
 
 
+// return actuator given mjsElement
+mjsActuator* mjs_asActuator(mjsElement* element) {
+  return element->elemtype == mjOBJ_ACTUATOR ? &(static_cast<mjCActuator*>(element)->spec) : nullptr;
+}
+
+
+
+// return sensor given mjsElement
+mjsSensor* mjs_asSensor(mjsElement* element) {
+  return element->elemtype == mjOBJ_SENSOR ? &(static_cast<mjCSensor*>(element)->spec) : nullptr;
+}
+
+
+
+// return flex given mjsElement
+mjsFlex* mjs_asFlex(mjsElement* element) {
+  return element->elemtype == mjOBJ_FLEX ? &(static_cast<mjCFlex*>(element)->spec) : nullptr;
+}
+
+
+
+// return pair given mjsElement
+mjsPair* mjs_asPair(mjsElement* element) {
+  return element->elemtype == mjOBJ_PAIR ? &(static_cast<mjCPair*>(element)->spec) : nullptr;
+}
+
+
+
+// return equality given mjsElement
+mjsEquality* mjs_asEquality(mjsElement* element) {
+  return element->elemtype == mjOBJ_EQUALITY ? &(static_cast<mjCEquality*>(element)->spec) : nullptr;
+}
+
+
+
+// return exclude given mjsElement
+mjsExclude* mjs_asExclude(mjsElement* element) {
+  return element->elemtype == mjOBJ_EXCLUDE ? &(static_cast<mjCBodyPair*>(element)->spec) : nullptr;
+}
+
+
+
+// return tendon given mjsElement
+mjsTendon* mjs_asTendon(mjsElement* element) {
+  return element->elemtype == mjOBJ_TENDON ? &(static_cast<mjCTendon*>(element)->spec) : nullptr;
+}
+
+
+
+// return numeric given mjsElement
+mjsNumeric* mjs_asNumeric(mjsElement* element) {
+  return element->elemtype == mjOBJ_NUMERIC ? &(static_cast<mjCNumeric*>(element)->spec) : nullptr;
+}
+
+
+
+// return text given mjsElement
+mjsText* mjs_asText(mjsElement* element) {
+  return element->elemtype == mjOBJ_TEXT ? &(static_cast<mjCText*>(element)->spec) : nullptr;
+}
+
+
+
+// return tuple given mjsElement
+mjsTuple* mjs_asTuple(mjsElement* element) {
+  return element->elemtype == mjOBJ_TUPLE ? &(static_cast<mjCTuple*>(element)->spec) : nullptr;
+}
+
+
+
+// return key given mjsElement
+mjsKey* mjs_asKey(mjsElement* element) {
+  return element->elemtype == mjOBJ_KEY ? &(static_cast<mjCKey*>(element)->spec) : nullptr;
+}
+
+
+
+// return mesh given mjsElement
+mjsMesh* mjs_asMesh(mjsElement* element) {
+  return element->elemtype == mjOBJ_MESH ? &(static_cast<mjCMesh*>(element)->spec) : nullptr;
+}
+
+
+
+// return hfield given mjsElement
+mjsHField* mjs_asHField(mjsElement* element) {
+  return element->elemtype == mjOBJ_HFIELD ? &(static_cast<mjCHField*>(element)->spec) : nullptr;
+}
+
+
+
+// return skin given mjsElement
+mjsSkin* mjs_asSkin(mjsElement* element) {
+  return element->elemtype == mjOBJ_SKIN ? &(static_cast<mjCSkin*>(element)->spec) : nullptr;
+}
+
+
+
+// return texture given mjsElement
+mjsTexture* mjs_asTexture(mjsElement* element) {
+  return element->elemtype == mjOBJ_TEXTURE ? &(static_cast<mjCTexture*>(element)->spec) : nullptr;
+}
+
+
+
+// return material given mjsElement
+mjsMaterial* mjs_asMaterial(mjsElement* element) {
+  return element->elemtype == mjOBJ_MATERIAL ? &(static_cast<mjCMaterial*>(element)->spec) : nullptr;
+}
+
+
+
 // set string
 void mjs_setString(mjString* dest, const char* text) {
   std::string* str = static_cast<std::string*>(dest);

src/user/user_api.h

@@ -218,6 +218,12 @@ 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);
 
+// Return spec's first element of selected type.
+MJAPI mjsElement* mjs_firstElement(mjSpec* s, mjtObj type);
+
+// Return spec's next element; return NULL if element is last.
+MJAPI mjsElement* mjs_nextElement(mjSpec* s, mjsElement* element);
+
 // Safely cast an element as mjsBody, or return NULL if the element is not an mjsBody.
 MJAPI mjsBody* mjs_asBody(mjsElement* element);
 
@@ -239,6 +245,54 @@ 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);
 
+// Safely cast an element as mjsActuator, or return NULL if the element is not an mjsActuator.
+MJAPI mjsActuator* mjs_asActuator(mjsElement* element);
+
+// Safely cast an element as mjsSensor, or return NULL if the element is not an mjsSensor.
+MJAPI mjsSensor* mjs_asSensor(mjsElement* element);
+
+// Safely cast an element as mjsFlex, or return NULL if the element is not an mjsFlex.
+MJAPI mjsFlex* mjs_asFlex(mjsElement* element);
+
+// Safely cast an element as mjsPair, or return NULL if the element is not an mjsPair.
+MJAPI mjsPair* mjs_asPair(mjsElement* element);
+
+// Safely cast an element as mjsEquality, or return NULL if the element is not an mjsEquality.
+MJAPI mjsEquality* mjs_asEquality(mjsElement* element);
+
+// Safely cast an element as mjsExclude, or return NULL if the element is not an mjsExclude.
+MJAPI mjsExclude* mjs_asExclude(mjsElement* element);
+
+// Safely cast an element as mjsTendon, or return NULL if the element is not an mjsTendon.
+MJAPI mjsTendon* mjs_asTendon(mjsElement* element);
+
+// Safely cast an element as mjsNumeric, or return NULL if the element is not an mjsNumeric.
+MJAPI mjsNumeric* mjs_asNumeric(mjsElement* element);
+
+// Safely cast an element as mjsText, or return NULL if the element is not an mjsText.
+MJAPI mjsText* mjs_asText(mjsElement* element);
+
+// Safely cast an element as mjsTuple, or return NULL if the element is not an mjsTuple.
+MJAPI mjsTuple* mjs_asTuple(mjsElement* element);
+
+// Safely cast an element as mjsKey, or return NULL if the element is not an mjsKey.
+MJAPI mjsKey* mjs_asKey(mjsElement* element);
+
+// Safely cast an element as mjsMesh, or return NULL if the element is not an mjsMesh.
+MJAPI mjsMesh* mjs_asMesh(mjsElement* element);
+
+// Safely cast an element as mjsHField, or return NULL if the element is not an mjsHField.
+MJAPI mjsHField* mjs_asHField(mjsElement* element);
+
+// Safely cast an element as mjsSkin, or return NULL if the element is not an mjsSkin.
+MJAPI mjsSkin* mjs_asSkin(mjsElement* element);
+
+// Safely cast an element as mjsTexture, or return NULL if the element is not an mjsTexture.
+MJAPI mjsTexture* mjs_asTexture(mjsElement* element);
+
+// Safely cast an element as mjsMaterial, or return NULL if the element is not an mjsMaterial.
+MJAPI mjsMaterial* mjs_asMaterial(mjsElement* element);
+
 
 //---------------------------------- Attribute setters ---------------------------------------------
 

src/user/user_model.cc

@@ -706,6 +706,70 @@ mjCBase* mjCModel::GetObject(mjtObj type, int id) {
 
 
 
+template <class T>
+static mjsElement* GetNext(std::vector<T*>& list, mjsElement* child) {
+  // TODO: use id for direct indexing instead of a loop
+  for (unsigned int i = 0; i < list.size()-1; i++) {
+    if (list[i]->spec.element == child) {
+      return list[i+1]->spec.element;
+    }
+  }
+  return nullptr;
+}
+
+
+
+// next object of specified type
+mjsElement* mjCModel::NextObject(mjsElement* object, mjtObj type) {
+  if (type == mjOBJ_UNKNOWN) {
+    if (!object) {
+      throw mjCError(NULL, "type must be specified if no element is given");
+    } else {
+      type = object->elemtype;
+    }
+  } else if (object && object->elemtype != type) {
+    throw mjCError(NULL, "element is not of requested type");
+  }
+
+  switch (type) {
+    case mjOBJ_ACTUATOR:
+      return object ? GetNext(actuators_, object) : actuators_[0];
+    case mjOBJ_SENSOR:
+      return object ? GetNext(sensors_, object) : sensors_[0];
+    case mjOBJ_FLEX:
+      return object ? GetNext(flexes_, object) : flexes_[0];
+    case mjOBJ_PAIR:
+      return object ? GetNext(pairs_, object) : pairs_[0];
+    case mjOBJ_EXCLUDE:
+      return object ? GetNext(excludes_, object) : excludes_[0];
+    case mjOBJ_EQUALITY:
+      return object ? GetNext(equalities_, object) : equalities_[0];
+    case mjOBJ_TENDON:
+      return object ? GetNext(tendons_, object) : tendons_[0];
+    case mjOBJ_NUMERIC:
+      return object ? GetNext(numerics_, object) : numerics_[0];
+    case mjOBJ_TEXT:
+      return object ? GetNext(texts_, object) : texts_[0];
+    case mjOBJ_TUPLE:
+      return object ? GetNext(tuples_, object) : tuples_[0];
+    case mjOBJ_KEY:
+      return object ? GetNext(keys_, object) : keys_[0];
+    case mjOBJ_MESH:
+      return object ? GetNext(meshes_, object) : meshes_[0];
+    case mjOBJ_HFIELD:
+      return object ? GetNext(hfields_, object) : hfields_[0];
+    case mjOBJ_SKIN:
+      return object ? GetNext(skins_, object) : skins_[0];
+    case mjOBJ_TEXTURE:
+      return object ? GetNext(textures_, object) : textures_[0];
+    case mjOBJ_MATERIAL:
+      return object ? GetNext(materials_, object) : materials_[0];
+    default:
+      return nullptr;
+  }
+}
+
+
 //------------------------ API FOR ACCESS TO PRIVATE VARIABLES -------------------------------------
 
 // compiled flag

src/user/user_model.h

@@ -203,6 +203,7 @@ class mjCModel : public mjCModel_, private mjSpec {
   // API for access to model elements (outside tree)
   int NumObjects(mjtObj type);              // number of objects in specified list
   mjCBase* GetObject(mjtObj type, int id);  // pointer to specified object
+  mjsElement* NextObject(mjsElement* object, mjtObj type = mjOBJ_UNKNOWN);  // next object of specified type
 
   // API for access to other variables
   bool IsCompiled() const;                                         // is model already compiled