Implement __getitem__ for Epoch

neo/core/epoch.py

@@ -24,10 +24,10 @@ def _new_epoch(cls, times=None, durations=None, labels=None, units=None,
                name=None, description=None, file_origin=None, annotations=None, segment=None):
     '''
     A function to map epoch.__new__ to function that
-    does not do the unit checking. This is needed for pickle to work. 
+    does not do the unit checking. This is needed for pickle to work.
     '''
-    e = Epoch(times=times, durations=durations, labels=labels, units=units, name=name, file_origin=file_origin,
-              description=description, **annotations)
+    e = Epoch(times=times, durations=durations, labels=labels, units=units, name=name,
+              file_origin=file_origin, description=description, **annotations)
     e.segment = segment
     return e
 
@@ -151,6 +151,16 @@ def __repr__(self):
                 label, time, dur in zip(labels, self.times, self.durations)]
         return '<Epoch: %s>' % ', '.join(objs)
 
+    def __getitem__(self, i):
+        '''
+        Get the item or slice :attr:`i`.
+        '''
+        obj = Epoch(times=super(Epoch, self).__getitem__(i))
+        obj._copy_data_complement(self)
+        obj.durations = self.durations[i]
+        obj.labels = self.labels[i]
+        return obj
+
     @property
     def times(self):
         return pq.Quantity(self)
@@ -232,10 +242,7 @@ def time_slice(self, t_start, t_stop):
             _t_stop = np.inf
 
         indices = (self >= _t_start) & (self <= _t_stop)
-
         new_epc = self[indices]
-        new_epc.durations = self.durations[indices]
-        new_epc.labels = self.labels[indices]
         return new_epc
 
     def as_array(self, units=None):

neo/test/coretest/test_epoch.py

@@ -497,6 +497,28 @@ def test_as_quantity(self):
         self.assertIsInstance(epc_as_q, pq.Quantity)
         assert_array_equal(times * pq.ms, epc_as_q)
 
+    def test_getitem(self):
+        times = [2, 3, 4, 5]
+        durations = [0.1, 0.2, 0.3, 0.4]
+        labels = ["A", "B", "C", "D"]
+        epc = Epoch(times * pq.ms, durations=durations * pq.ms, labels=labels)
+        single_epoch = epc[2]
+        self.assertIsInstance(single_epoch, Epoch)
+        assert_array_equal(single_epoch.times, np.array([4.0]))
+        assert_array_equal(single_epoch.durations, np.array([0.3]))
+        assert_array_equal(single_epoch.labels, np.array(["C"]))
+
+    def test_slice(self):
+        times = [2, 3, 4, 5]
+        durations = [0.1, 0.2, 0.3, 0.4]
+        labels = ["A", "B", "C", "D"]
+        epc = Epoch(times * pq.ms, durations=durations * pq.ms, labels=labels)
+        single_epoch = epc[1:3]
+        self.assertIsInstance(single_epoch, Epoch)
+        assert_array_equal(single_epoch.times, np.array([3.0, 4.0]))
+        assert_array_equal(single_epoch.durations, np.array([0.2, 0.3]))
+        assert_array_equal(single_epoch.labels, np.array(["B", "C"]))
+
 
 class TestDuplicateWithNewData(unittest.TestCase):
     def setUp(self):