Skip to content

Commit f0f4cc6

Browse files
authored
Merge pull request #1539 from zm711/slice-none
Add test that all rawios accept `slice(None)`
2 parents 3b5488c + e2c1cf8 commit f0f4cc6

File tree

2 files changed

+35
-3
lines changed

2 files changed

+35
-3
lines changed

neo/rawio/medrawio.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -240,9 +240,18 @@ def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, strea
240240
self.sess.set_channel_active(self._stream_info[stream_index]["raw_chans"])
241241
num_channels = len(self._stream_info[stream_index]["raw_chans"])
242242
self.sess.set_reference_channel(self._stream_info[stream_index]["raw_chans"][0])
243+
244+
# in the case we have a slice or we give an ArrayLike we need to iterate through the channels
245+
# in order to activate them.
243246
else:
244-
if any(channel_indexes < 0):
245-
raise IndexError(f"Can not index negative channels: {channel_indexes}")
247+
if isinstance(channel_indexes, slice):
248+
start = channel_indexes.start or 0
249+
stop = channel_indexes.stop or len(self._stream_info[stream_index]["raw_chans"])
250+
step = channel_indexes.step or 1
251+
channel_indexes = [ch for ch in range(start, stop, step)]
252+
else:
253+
if any(channel_indexes < 0):
254+
raise IndexError(f"Can not index negative channels: {channel_indexes}")
246255
# Set all channels to be inactive, then selectively set some of them to be active
247256
self.sess.set_channel_inactive("all")
248257
for i, channel_idx in enumerate(channel_indexes):

neo/test/rawiotest/rawio_compliance.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,7 @@ def read_analogsignals(reader):
172172
channel_names = signal_channels["name"][mask]
173173
channel_ids = signal_channels["id"][mask]
174174

175-
# acces by channel inde/ids/names should give the same chunk
175+
# acces by channel index/ids/names should give the same chunk
176176
channel_indexes2 = channel_indexes[::2]
177177
channel_names2 = channel_names[::2]
178178
channel_ids2 = channel_ids[::2]
@@ -214,6 +214,29 @@ def read_analogsignals(reader):
214214
)
215215
np.testing.assert_array_equal(raw_chunk0, raw_chunk1)
216216

217+
# test slice(None). This should return the same array as giving
218+
# all channel indexes or using None as an argument in `get_analogsignal_chunk`
219+
# see https://github.com/NeuralEnsemble/python-neo/issues/1533
220+
221+
raw_chunk_slice_none = reader.get_analogsignal_chunk(
222+
block_index=block_index,
223+
seg_index=seg_index,
224+
i_start=i_start,
225+
i_stop=i_stop,
226+
stream_index=stream_index,
227+
channel_indexes=slice(None)
228+
)
229+
raw_chunk_channel_indexes = reader.get_analogsignal_chunk(
230+
block_index=block_index,
231+
seg_index=seg_index,
232+
i_start=i_start,
233+
i_stop=i_stop,
234+
stream_index=stream_index,
235+
channel_indexes=channel_indexes
236+
)
237+
238+
np.testing.assert_array_equal(raw_chunk_slice_none, raw_chunk_channel_indexes)
239+
217240
# test prefer_slice=True/False
218241
if nb_chan >= 3:
219242
for prefer_slice in (True, False):

0 commit comments

Comments
 (0)