Skip to content

Commit

Permalink
Support iterating over TdmsFile and TdmsGroup
Browse files Browse the repository at this point in the history
Because these act like dictionaries, iterating over them
produces group and channel names.
  • Loading branch information
adamreeve committed Mar 29, 2020
1 parent f1974e7 commit 5af8c8a
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 2 deletions.
6 changes: 6 additions & 0 deletions nptdms/tdms.py
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,9 @@ def close(self):
self._reader.close()
self._reader = None

def __iter__(self):
return iter(self._groups)

def __getitem__(self, group_name):
""" Retrieve a TDMS group from the file by name
"""
Expand Down Expand Up @@ -426,6 +429,9 @@ def as_dataframe(self, time_index=False, absolute_time=False, scaled_data=True):

return pandas_export.from_group(self, time_index, absolute_time, scaled_data)

def __iter__(self):
return iter(self._channels)

def __getitem__(self, channel_name):
""" Retrieve a TDMS channel from this group by name
"""
Expand Down
19 changes: 17 additions & 2 deletions nptdms/test/test_tdms_file.py
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,22 @@ def test_iterate_channel_data_in_read_mode():
compare_arrays(actual_data, expected_channel_data)


def test_invalid_offset_throws():
def test_iterate_file_and_groups():
""" Test iterating over TdmsFile and TdmsGroup uses key values
"""
test_file, expected_data = scenarios.chunked_segment().values

with test_file.get_tempfile() as temp_file:
tdms_file = TdmsFile.read(temp_file.file)
for group_name in tdms_file:
group = tdms_file[group_name]
for channel_name in group:
channel = group[channel_name]
expected_channel_data = expected_data[(group_name, channel_name)]
compare_arrays(channel.data, expected_channel_data)


def test_invalid_offset_in_read_data_throws():
""" Exception is thrown when reading a subset of data with an invalid offset
"""
test_file, expected_data = scenarios.single_segment_with_one_channel().values
Expand All @@ -230,7 +245,7 @@ def test_invalid_offset_throws():
assert "offset must be non-negative" in str(exc_info.value)


def test_invalid_length_throws():
def test_invalid_length_in_read_data_throws():
""" Exception is thrown when reading a subset of data with an invalid length
"""
test_file, expected_data = scenarios.single_segment_with_one_channel().values
Expand Down

0 comments on commit 5af8c8a

Please sign in to comment.