Merge pull request #1307 from kjohnsen/segment-add

add `add()` function to `Container` for use on `Segment` and `Block`

neo/core/container.py

@@ -327,6 +327,25 @@ def size(self):
         return {name: len(getattr(self, name))
                     for name in self._child_containers}
 
+    @property
+    def _container_lookup(self):
+        return {
+            cls_name: getattr(self, container_name)
+            for cls_name, container_name in zip(self._child_objects, self._child_containers)
+        }
+
+    def _get_container(self, cls):
+        if hasattr(cls, "proxy_for"):
+            cls = cls.proxy_for
+        return self._container_lookup[cls.__name__]
+
+    def add(self, *objects):
+        """Add a new Neo object to the Container"""
+        for obj in objects:
+            container = self._get_container(obj.__class__)
+            container.append(obj)
+
+
     def filter(self, targdict=None, data=True, container=False, recursive=True,
                objects=None, **kwargs):
         """

neo/core/group.py

@@ -134,18 +134,6 @@ def __init__(self, objects=None, name=None, description=None, file_origin=None,
         doc="list of Groups contained in this group"
     )
 
-    @property
-    def _container_lookup(self):
-        return {
-            cls_name: getattr(self, container_name)
-            for cls_name, container_name in zip(self._child_objects, self._child_containers)
-        }
-
-    def _get_container(self, cls):
-        if hasattr(cls, "proxy_for"):
-            cls = cls.proxy_for
-        return self._container_lookup[cls.__name__]
-
     def add(self, *objects):
         """Add a new Neo object to the Group"""
         for obj in objects:

neo/test/coretest/test_block.py

@@ -466,6 +466,33 @@ def test_segment_list(self):
         blk.segments += [Segment(), Segment()]
         assert len(blk.segments) == 3
 
+    def test_add(self):
+        blk = self.blocks[0]
+        new_blk = simple_block()
+        n_groups_start = len(new_blk.groups)
+        for group in blk.groups:
+            assert group not in new_blk.groups
+            new_blk.add(group)
+            assert group in new_blk.groups
+        assert len(new_blk.groups) == n_groups_start + len(blk.groups)
+
+        n_segs_start = len(new_blk.segments)
+        for seg in blk.segments:
+            assert seg not in new_blk.segments
+            new_blk.add(seg)
+            assert seg in new_blk.segments
+        assert len(new_blk.segments) == n_segs_start + len(blk.segments)
+
+        # test adding multiple at once
+        blk = self.blocks[1]
+        n_groups_start = len(new_blk.groups)
+        new_blk.add(*blk.groups)
+        assert len(new_blk.groups) == n_groups_start + len(blk.groups)
+
+        n_segs_start = len(new_blk.segments)
+        new_blk.add(*blk.segments)
+        assert len(new_blk.segments) == n_segs_start + len(blk.segments)
+
 
 if __name__ == "__main__":
     unittest.main()

neo/test/coretest/test_segment.py

@@ -642,6 +642,33 @@ def test__deepcopy(self):
                 for child in getattr(seg1_copy, childtype, []):
                     self.assertEqual(id(child.segment), id(seg1_copy))
 
+    def test_add(self):
+        seg = Segment()
+
+        reader = ExampleRawIO(filename='my_filename.fake')
+        reader.parse_header()
+
+        proxy_anasig = AnalogSignalProxy(rawio=reader,
+                                        stream_index=0, inner_stream_channels=None,
+                                        block_index=0, seg_index=0)
+        seg.add(proxy_anasig)
+        assert len(seg.analogsignals) == 1
+
+        proxy_st = SpikeTrainProxy(rawio=reader, spike_channel_index=0,
+                                   block_index=0, seg_index=0)
+        seg.add(proxy_st)
+        assert len(seg.spiketrains) == 1
+
+        proxy_event = EventProxy(rawio=reader, event_channel_index=0,
+                                 block_index=0, seg_index=0)
+        seg.add(proxy_event)
+        assert len(seg.events) == 1
+
+        proxy_epoch = EpochProxy(rawio=reader, event_channel_index=1,
+                                 block_index=0, seg_index=0)
+        seg.add(proxy_epoch)
+        assert len(seg.epochs) == 1
+
 
 if __name__ == "__main__":
     unittest.main()