Allow filtering PyMJCF find_all results by default class

dm_control/mjcf/element.py

@@ -453,7 +453,8 @@ def find(self, namespace, identifier):
     return found_element
 
   def find_all(self, namespace,
-               immediate_children_only=False, exclude_attachments=False):
+               immediate_children_only=False, exclude_attachments=False,
+               dclass=None):
     """Finds all elements of a particular kind.
 
     The `namespace` argument specifies the kind of element to
@@ -472,6 +473,9 @@ def find_all(self, namespace,
       exclude_attachments: (optional) A boolean, if `True` then elements
         belonging to attached models are excluded.
 
+      dclass: (optional) A default class name string to match. If None,
+        elements of all default classes are returned.
+
     Returns:
       A list of `mjcf.Element`.
 
@@ -491,10 +495,15 @@ def find_all(self, namespace,
               namespace + constants.NAMESPACE_SEPARATOR) or
           # Attachment frames are considered part of the "body" namespace.
           namespace == constants.BODY and isinstance(child, _AttachmentFrame)):
-        out.append(child)
+        default = getattr(child, 'dclass', None)
+        # A PyMJCF dclass is an Element, rather than the XML class string.
+        class_name = getattr(default, 'dclass', default)
+        if dclass is None or class_name == dclass:
+          out.append(child)
       if not immediate_children_only:
-        out.extend(child.find_all(namespace,
-                                  exclude_attachments=exclude_attachments))
+        out.extend(child.find_all(
+            namespace, exclude_attachments=exclude_attachments,
+            dclass=dclass))
     return out
 
   def enter_scope(self, scope_identifier):

dm_control/mjcf/element_test.py

@@ -776,6 +776,23 @@ def testFindAll(self):
     self.assertLen(b_0.find_all('joint', immediate_children_only=True), 1)
     self.assertLen(b_0.find_all('joint', exclude_attachments=True), 2)
 
+  def testFindAllFiltersByDefaultClass(self):
+    root = element.RootElement(model='filter_defaults')
+    default = root.default.add('default', dclass='red_sites')
+    body = root.worldbody.add('body', name='body')
+    red = body.add('site', name='red', dclass=default)
+    plain = body.add('site', name='plain')
+    nested = body.add('body', name='nested')
+    red_nested = nested.add('site', name='red_nested', dclass=default)
+
+    self.assertEqual(
+        root.find_all('site', dclass='red_sites'), [red, red_nested])
+    self.assertEqual(root.find_all('site'), [red, plain, red_nested])
+    self.assertEqual(root.find_all('site', dclass='not_used'), [])
+    self.assertEqual(
+        body.find_all('site', immediate_children_only=True,
+                      dclass='red_sites'), [red])
+
   def testFindAllFrameJoints(self):
     root_model = parser.from_path(_TEST_MODEL_XML)
     root_model.model = 'model'