Add method to find objects of a given neurodata type

CHANGELOG.md

@@ -6,6 +6,7 @@
 - Lifted the `<1.11` cap on the `linkml` and `linkml-runtime` termset extras and bumped the `dandi` docs dependency to `>=0.76.5`. dandi 0.76.5 lifts its `click<8.2` bound ([dandi/dandi-cli#1883](https://github.com/dandi/dandi-cli/pull/1883)), which had conflicted with the `click>=8.2` requirement of linkml 1.11. @rly [#2217](https://github.com/NeurodataWithoutBorders/pynwb/pull/2217)
 
 ### 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 remote-read support to `pynwb.read_nwb`. The function now accepts remote URLs (`s3://`, `gs://`, `abfs://`, `https://`, etc.) and dispatches to the right backend based on the URL: `.zarr` suffixes (and DANDI Zarr assets under `/zarr/`) are read with `NWBZarrIO`, everything else with `NWBHDF5IO`. Remote files are opened through `fsspec`, which now uses the URL's actual scheme instead of the previous hardcoded `fsspec.filesystem("http")` that mishandled non-HTTP schemes. @h-mayorquin [#2190](https://github.com/NeurodataWithoutBorders/pynwb/pull/2190)
 
 ### Fixed

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