Fix reading files with truncated metadata in the final segment (#273)

nptdms/reader.py

@@ -245,11 +245,11 @@ def _read_segment_metadata(
 
     def _read_lead_in(self, file, segment_position, is_index_file=False):
         lead_in_bytes = file.read(28)
+        if len(lead_in_bytes) < 28:
+            raise EOFError
 
         expected_tag = b'TDSh' if is_index_file else b'TDSm'
         tag = lead_in_bytes[:4]
-        if tag == b'':
-            raise EOFError
         if tag != expected_tag:
             raise ValueError(
                 "Segment does not start with %r, but with %r" % (expected_tag, tag))
@@ -281,11 +281,16 @@ def _read_lead_in(self, file, segment_position, is_index_file=False):
         segment_incomplete = next_segment_offset == 0xFFFFFFFFFFFFFFFF
         if segment_incomplete:
             # Segment size is unknown. This can happen if LabVIEW crashes.
-            # Try to read until the end of the file.
+            next_segment_pos = self._get_data_file_size()
+            if next_segment_pos < data_position:
+                # Metadata wasn't completely written and don't have any data in this segment,
+                # don't try to read any metadata
+                log.warning("Last segment metadata is incomplete")
+                raise EOFError
+            # Try to read until the end of the file if we have complete metadata
             log.warning(
                 "Last segment of file has unknown size, "
                 "will attempt to read to the end of the file")
-            next_segment_pos = self._get_data_file_size()
         else:
             log.debug("Next segment offset = %d, raw data offset = %d, data size = %d b",
                       next_segment_offset, raw_data_offset, next_segment_offset - raw_data_offset)

nptdms/test/test_tdms_file.py

@@ -770,6 +770,53 @@ def test_incomplete_segment_with_string_data():
     assert len(channel) == 0
 
 
+def test_truncated_metadata_in_last_segment():
+    """ Test the scenario where writing the file was aborted with part of the metadata written
+    """
+    test_file = GeneratedFile()
+    test_file.add_segment(
+        ("kTocMetaData", "kTocRawData", "kTocNewObjList"),
+        segment_objects_metadata(
+            channel_metadata("/'group'/'channel1'", 3, 2),
+            channel_metadata("/'group'/'channel2'", 3, 2),
+        ),
+        "01 00 00 00" "02 00 00 00"
+        "03 00 00 00" "04 00 00 00"
+    )
+    first_segment_size = len(test_file)
+    test_file.add_segment(
+        ("kTocMetaData", "kTocRawData", "kTocNewObjList"),
+        segment_objects_metadata(
+            channel_metadata("/'group'/'channel1'", 3, 2),
+            channel_metadata("/'group'/'channel2'", 3, 2),
+        ),
+        "",
+        incomplete=True
+    )
+    total_size = len(test_file)
+    expected_data = {
+        ('group', 'channel1'): np.array([1, 2], dtype=np.int32),
+        ('group', 'channel2'): np.array([3, 4], dtype=np.int32),
+    }
+
+    with test_file.get_tempfile() as temp_file:
+        with tempfile.NamedTemporaryFile(suffix=".tdms") as truncated_file:
+            all_data = temp_file.file.read()
+            for end_point in range(first_segment_size + 1, total_size - 1):
+                truncated_file.file.seek(0)
+                truncated_file.file.write(all_data[:end_point])
+                truncated_file.file.seek(0)
+                try:
+                    tdms_data = TdmsFile.read(truncated_file.file)
+                except Exception as exc:
+                    raise RuntimeError(
+                        "Failed to read file with segment truncated after "
+                        "%d bytes: %s" % (end_point - first_segment_size, exc))
+                for ((group, channel), expected_values) in expected_data.items():
+                    channel_obj = tdms_data[group][channel]
+                    compare_arrays(channel_obj.data, expected_values)
+
+
 def test_slash_and_space_in_name():
     """Test name like '01/02/03 something'"""
 

nptdms/test/util.py

@@ -300,6 +300,9 @@ def get_bytes_io_file(self):
         file.seek(0)
         return file
 
+    def __len__(self):
+        return sum(len(s[0]) + len(s[1]) + len(s[2]) for s in self._content)
+
     def _get_contents(self):
         contents = b''
         for segment in self._content: