Add method to find objects of a given neurodata type

CHANGELOG.md

@@ -8,6 +8,7 @@
 - Added `EventsTable` examples to the NWB file basics and behavior tutorials. @rly [#2156](https://github.com/NeurodataWithoutBorders/pynwb/pull/2156)
 
 ### Added
+- Added `collect_descendants_of_type` method to `NWBContainer` for finding all descendant objects of a given neurodata type (by class or by string name), addressing the long-standing request in #560. @h-mayorquin [#2189](https://github.com/NeurodataWithoutBorders/pynwb/pull/2189)
 - Added support for NWB Schema 2.10.0 ([NWBEP001](https://nwb-schema.readthedocs.io/)), which introduces the `EventsTable`, `TimestampVectorData`, and `DurationVectorData` neurodata types and a new `events` group on `NWBFile` for storing `EventsTable` instances. Use `NWBFile.add_events_table()` to add an `EventsTable` and `NWBFile.get_events_table()` to retrieve one. NWB Schema 2.10.0 also incorporates hdmf-common-schema 1.9.0, which adds the `MeaningsTable` neurodata type (re-exported via `hdmf.common`) and support for attaching one or more `MeaningsTable` instances to a `DynamicTable` to document the meanings of values in a column. See the [NWB Schema release notes](https://nwb-schema.readthedocs.io/en/latest/format_release_notes.html) and the [hdmf-common-schema release notes](https://hdmf-common-schema.readthedocs.io/en/latest/format_release_notes.html) for the full list of changes. @rly [#2156](https://github.com/NeurodataWithoutBorders/pynwb/pull/2156)
 - Added support for HERD (HDMF External Resources Data Structure) as the `external_resources` field on `NWBFile`, enabling users to associate external resource annotations (e.g., ontology term mappings) with their NWB files. `link_resources` and `get_external_resources` are inherited from `HERDManager` in hdmf. @mavaylon1, @rly [#2111](https://github.com/NeurodataWithoutBorders/pynwb/pull/2111)
 - Added `get_starting_time()` and `get_duration()` methods to `TimeSeries` to get the starting time and duration of the time series. @h-mayorquin [#2146](https://github.com/NeurodataWithoutBorders/pynwb/pull/2146)

src/pynwb/core.py

@@ -85,7 +85,24 @@ def data_type(self):
 @register_class('NWBContainer', CORE_NAMESPACE)
 class NWBContainer(NWBMixin, Container):
 
-    pass
+    @docval({'name': 'neurodata_type', 'type': (type, str),
+             'doc': ('A PyNWB container class, or the string name of a neurodata_type. '
+                     'Subclass matches are included. For type names that may be ambiguous '
+                     'across loaded namespaces, pass the class directly.')})
+    def collect_descendants_of_type(self, **kwargs):
+        """Return all descendants that are instances of the given neurodata type.
+
+        Walks the entire subtree of this container and returns every descendant that is an
+        instance of the given class. Subclasses are included. ``self`` is excluded from the
+        result. Returns an empty list if there are no matches.
+        """
+        neurodata_type = kwargs['neurodata_type']
+        if isinstance(neurodata_type, str):
+            cls = self._get_type_map().get_dt_container_cls(neurodata_type)
+        else:
+            cls = neurodata_type
+        return [obj for obj in self.all_children()
+                if obj is not self and isinstance(obj, cls)]
 
 
 @register_class('NWBDataInterface', CORE_NAMESPACE)

tests/unit/test_file.py

@@ -16,6 +16,7 @@
 from pynwb.file import Subject, _add_missing_timezone
 from pynwb.epoch import TimeIntervals
 from pynwb.ecephys import ElectricalSeries, ElectrodesTable
+from pynwb.behavior import SpatialSeries
 from pynwb.testing import TestCase, remove_test_file
 
 
@@ -443,6 +444,63 @@ def test_all_children(self):
         self.assertIn(device, children)
         self.assertIn(elecgrp, children)
 
+    def test_collect_descendants_of_type_class_arg(self):
+        """Passing a class returns every descendant that is an instance of it."""
+        ts1 = TimeSeries(name='ts1', data=[0, 1, 2], unit='grams', timestamps=[0.0, 0.1, 0.2])
+        ts2 = TimeSeries(name='ts2', data=[3, 4, 5], unit='grams', timestamps=[0.3, 0.4, 0.5])
+        self.nwbfile.add_acquisition(ts1)
+        self.nwbfile.add_acquisition(ts2)
+        result = self.nwbfile.collect_descendants_of_type(TimeSeries)
+        self.assertIn(ts1, result)
+        self.assertIn(ts2, result)
+        self.assertEqual(len(result), 2)
+
+    def test_collect_descendants_of_type_subclass_inclusion(self):
+        """Querying with a base class returns subclass instances too."""
+        ts = TimeSeries(name='ts', data=[0, 1, 2], unit='grams', timestamps=[0.0, 0.1, 0.2])
+        ss = SpatialSeries(name='ss', data=[[0, 0], [1, 1], [2, 2]],
+                           reference_frame='origin', timestamps=[0.0, 0.1, 0.2])
+        self.nwbfile.add_acquisition(ts)
+        self.nwbfile.add_acquisition(ss)
+        all_ts = self.nwbfile.collect_descendants_of_type(TimeSeries)
+        self.assertIn(ts, all_ts)
+        self.assertIn(ss, all_ts)
+        only_ss = self.nwbfile.collect_descendants_of_type(SpatialSeries)
+        self.assertEqual(only_ss, [ss])
+
+    def test_collect_descendants_of_type_string_arg(self):
+        """Passing a string neurodata_type name resolves to the same class."""
+        ts = TimeSeries(name='ts', data=[0, 1, 2], unit='grams', timestamps=[0.0, 0.1, 0.2])
+        self.nwbfile.add_acquisition(ts)
+        by_class = self.nwbfile.collect_descendants_of_type(TimeSeries)
+        by_string = self.nwbfile.collect_descendants_of_type('TimeSeries')
+        self.assertEqual(by_class, by_string)
+        self.assertIn(ts, by_string)
+
+    def test_collect_descendants_of_type_no_matches(self):
+        """Querying for a type with zero instances returns an empty list."""
+        result = self.nwbfile.collect_descendants_of_type(SpatialSeries)
+        self.assertEqual(result, [])
+
+    def test_collect_descendants_of_type_excludes_self(self):
+        """A container does not include itself in its own descendant search."""
+        ts = TimeSeries(name='ts', data=[0, 1, 2], unit='grams', timestamps=[0.0, 0.1, 0.2])
+        self.nwbfile.add_acquisition(ts)
+        result = ts.collect_descendants_of_type(TimeSeries)
+        self.assertNotIn(ts, result)
+
+    def test_collect_descendants_of_type_scoped_to_subtree(self):
+        """Calling on a sub-container only searches its subtree, not siblings."""
+        ts_a = TimeSeries(name='ts_a', data=[0, 1, 2], unit='grams', timestamps=[0.0, 0.1, 0.2])
+        ts_b = TimeSeries(name='ts_b', data=[3, 4, 5], unit='grams', timestamps=[0.3, 0.4, 0.5])
+        mod_a = self.nwbfile.create_processing_module(name='mod_a', description='a')
+        mod_b = self.nwbfile.create_processing_module(name='mod_b', description='b')
+        mod_a.add(ts_a)
+        mod_b.add(ts_b)
+        in_a = mod_a.collect_descendants_of_type(TimeSeries)
+        self.assertEqual(in_a, [ts_a])
+        self.assertNotIn(ts_b, in_a)
+
     def test_fail_if_source_script_file_name_without_source_script(self):
         with self.assertRaises(ValueError):
             # <-- source_script_file_name without source_script is not allowed