Add option to specify the TDMS format version to write (#265)

nptdms/reader.py

@@ -36,6 +36,7 @@ def __init__(self, tdms_file):
         self._file_path = None
         self._index_file_path = None
         self._segment_channel_offsets = {}
+        self.tdms_version = None
 
         if hasattr(tdms_file, "read"):
             # Is a file
@@ -267,8 +268,12 @@ def _read_lead_in(self, file, segment_position, is_index_file=False):
         # Next four bytes are version number, then 8 bytes each for the offset values
         (version, next_segment_offset, raw_data_offset) = _struct_unpack(endianness + 'lQQ', lead_in_bytes[8:28])
 
-        if version not in (4712, 4713):
-            log.warning("Unrecognised version number.")
+        if self.tdms_version is None:
+            if version not in (4712, 4713):
+                log.warning("Unrecognised version number: %d" % version)
+            self.tdms_version = version
+        elif self.tdms_version != version:
+            log.warning("Segment version mismatch, %d != %d" % (version, self.tdms_version))
 
         # Calculate data and next segment position
         lead_size = 7 * 4

nptdms/tdms.py

@@ -123,6 +123,7 @@ def __init__(self, file, raw_timestamps=False, memmap_dir=None, read_metadata_on
         self._groups = OrderedDict()
         self._properties = OrderedDict()
         self._channel_data = {}
+        self._tdms_version = 0
         self.data_read = False
 
         self._reader = TdmsReader(file)
@@ -140,6 +141,12 @@ def groups(self):
 
         return list(self._groups.values())
 
+    @property
+    def tdms_version(self):
+        """ The TDMS format version of this file
+        """
+        return self._tdms_version
+
     @property
     def properties(self):
         """ Return the properties of this file as a dictionary
@@ -229,6 +236,7 @@ def __exit__(self, exc_type, exc_value, traceback):
 
     def _read_file(self, tdms_reader, read_metadata_only, keep_open):
         tdms_reader.read_metadata(require_segment_indexes=keep_open)
+        self._tdms_version = tdms_reader.tdms_version
 
         # Use object metadata to build group and channel objects
         group_properties = OrderedDict()

nptdms/test/test_tdms_file.py

@@ -1056,3 +1056,26 @@ def test_channel_name_autocompletion():
     completion_options = group._ipython_key_completions_()
     assert "Channel1" in completion_options
     assert "Channel2" in completion_options
+
+
+def test_warning_on_unrecognized_version(caplog):
+    test_file = GeneratedFile()
+    test_file.add_segment(*basic_segment(), version=4714)
+
+    _ = test_file.load()
+
+    assert "WARNING" in caplog.text
+    assert "Unrecognised version number" in caplog.text
+    assert "4714" in caplog.text
+
+
+def test_warning_on_version_mismatch(caplog):
+    test_file = GeneratedFile()
+    test_file.add_segment(*basic_segment(), version=4713)
+    test_file.add_segment(*basic_segment(), version=4712)
+
+    _ = test_file.load()
+
+    assert "WARNING" in caplog.text
+    assert "Segment version mismatch" in caplog.text
+    assert "4712 != 4713" in caplog.text

nptdms/test/util.py

@@ -230,7 +230,7 @@ class GeneratedFile(object):
     def __init__(self):
         self._content = []
 
-    def add_segment(self, toc, metadata, data, incomplete=False, binary_data=False):
+    def add_segment(self, toc, metadata, data, incomplete=False, binary_data=False, version=4713):
         metadata_bytes = _hex_to_bytes(metadata)
         data_bytes = data if binary_data else _hex_to_bytes(data)
         if toc is not None:
@@ -252,7 +252,7 @@ def add_segment(self, toc, metadata, data, incomplete=False, binary_data=False):
                 else:
                     raise ValueError("Unrecognised TOC value: %s" % toc_item)
             lead_in += struct.pack('<i', toc_mask)
-            lead_in += _hex_to_bytes("69 12 00 00")
+            lead_in += struct.pack('<l', version)
             next_segment_offset = len(metadata_bytes) + len(data_bytes)
             raw_data_offset = len(metadata_bytes)
             if incomplete:

nptdms/test/writer/test_acceptance_tests.py

@@ -4,6 +4,7 @@
 from io import BytesIO
 import numpy as np
 import os
+import pytest
 import tempfile
 
 from nptdms import TdmsFile, TdmsWriter, RootObject, GroupObject, ChannelObject, types
@@ -26,6 +27,7 @@ def test_can_read_tdms_file_after_writing():
     a_output = tdms_file["group"]["a"]
     b_output = tdms_file["group"]["b"]
 
+    assert tdms_file.tdms_version == 4712
     assert a_output.data_type == types.SingleFloat
     assert b_output.data_type == types.DoubleFloat
     assert a_output.dtype == np.dtype('float32')
@@ -377,3 +379,27 @@ def test_can_write_complex():
     assert len(output_data) == 2
     assert output_data[0] == input_complex128_data[0]
     assert output_data[1] == input_complex128_data[1]
+
+
+def test_specifying_version():
+    data = np.linspace(0.0, 1.0, 100, dtype='float32')
+    channel = ChannelObject("group", "a", data)
+
+    output_file = BytesIO()
+    with TdmsWriter(output_file, version=4713) as tdms_writer:
+        tdms_writer.write_segment([channel])
+
+    output_file.seek(0)
+    tdms_file = TdmsFile(output_file)
+
+    assert tdms_file.tdms_version == 4713
+
+
+def test_specifying_invalid_version():
+    output_file = BytesIO()
+
+    with pytest.raises(ValueError) as exception:
+        _ = TdmsWriter(output_file, version=4714)
+    error_message = str(exception.value)
+
+    assert "4712,4713" in error_message

nptdms/writer.py

@@ -28,7 +28,7 @@ class TdmsWriter(object):
             tdms_writer.write_segment(segment_data)
     """
 
-    def __init__(self, file, mode='w'):
+    def __init__(self, file, mode='w', version=4712):
         """Initialise a new TDMS writer
 
         :param file: Either the path to the tdms file to open or an already
@@ -37,10 +37,18 @@ def __init__(self, file, mode='w'):
             This will be passed through to Python's ``open`` function with 'b' appended
             to ensure the file is opened in binary mode.
             For example, use 'w' (the default) to open a new file or 'a' to append to an existing TDMS file.
+        :param version: The TDMS format version to write, which must be either 4712 (the default) or 4713.
+            It's important that if you are appending segments to an
+            existing TDMS file, this matches the existing file version (this can be queried with the
+            :py:attr:`~nptdms.TdmsFile.tdms_version` property).
         """
+        valid_versions = (4712, 4713)
+        if version not in valid_versions:
+            raise ValueError("version must be one of %s" % ",".join("%d" % v for v in valid_versions))
         self._file = None
         self._file_path = None
         self._file_mode = mode
+        self._tdms_version = version
 
         if hasattr(file, "read"):
             # Is a file
@@ -62,7 +70,7 @@ def write_segment(self, objects):
 
         :param objects: A list of TdmsObject instances to write
         """
-        segment = TdmsSegment(objects)
+        segment = TdmsSegment(objects, version=self._tdms_version)
         segment.write(self._file)
 
     def __enter__(self):
@@ -77,16 +85,18 @@ class TdmsSegment(object):
     """A segment of data to be written to a file
     """
 
-    def __init__(self, objects):
+    def __init__(self, objects, version=4712):
         """Initialise a new segment of TDMS data
 
         :param objects: A list of TdmsObject instances.
+        :param version: The TDMS format version to write, which must be either 4712 (the default) or 4713.
         """
         paths = set(obj.path for obj in objects)
         if len(paths) != len(objects):
             raise ValueError("Duplicate object paths found")
 
         self.objects = objects
+        self._tdms_version = version
 
     def write(self, file):
         metadata = self.metadata()
@@ -139,8 +149,7 @@ def leadin(self, toc, metadata_size):
             toc_mask = toc_mask | toc_properties[toc_flag]
         leadin.append(Int32(toc_mask))
 
-        tdms_version = 4712
-        leadin.append(Int32(tdms_version))
+        leadin.append(Int32(self._tdms_version))
 
         next_segment_offset = metadata_size + self._data_size()
         raw_data_offset = metadata_size