feature: raw_timestamps support for writer and defragment (#320)

nptdms/test/test_example_files.py

@@ -1,10 +1,12 @@
-""" Test reading example TDMS files
-"""
+""" Test reading example TDMS files """
 
 import os
+from io import BytesIO
+
 import numpy as np
-from nptdms import tdms
+import pytest
 
+from nptdms import tdms, TdmsWriter, TdmsFile
 
 DATA_DIR = os.path.dirname(os.path.realpath(__file__)) + '/data'
 
@@ -33,6 +35,49 @@ def test_raw_format():
                                     0.2517777])
 
 
+def tdms_files_assert_equal(tdms1: TdmsFile, tdms2: TdmsFile):
+    """Assert that two TdmsFile instances are equal"""
+    # verify file properties
+    for p in tdms1.properties:
+        np.testing.assert_equal(tdms1.properties[p], tdms2.properties[p])
+    # verify group
+    for group in tdms1.groups():
+        # verify group properties
+        for p in group.properties:
+            np.testing.assert_equal(group.properties[p], tdms2[group.name].properties[p])
+        # verify channels
+        for channel in group.channels():
+            # verify channel data
+            np.testing.assert_equal(channel.data, tdms2[group.name][channel.name].data)
+            # verify channel properties
+            for p in channel.properties:
+                np.testing.assert_equal(channel.properties[p], tdms2[group.name][channel.name].properties[p])
+
+
+@pytest.mark.parametrize("tdms_file", [
+    'raw_timestamps.tdms',  # Test defragmentation of a file with raw timestamps
+    # 'raw1.tdms', # <- cannot defragment this file (ValueError: Channel data must be a 1d array)
+    'Digital_Input.tdms',
+    'big_endian.tdms',
+])
+def test_defragment(tdms_file):
+    """Test defragmentation round trip for a TDMS file"""
+    test_file_path = DATA_DIR + '/' + tdms_file
+    output_file = BytesIO()
+
+    # verify we can defragment a file with raw timestamps
+    TdmsWriter.defragment(test_file_path, output_file)
+
+    # rewind output file BytesIO instance, so it can read it back in as a TdmsFile
+    output_file.seek(0)
+
+    # verify that both TdmsFile objects are the same
+    tdms_files_assert_equal(
+        tdms.TdmsFile(test_file_path, raw_timestamps=True),
+        tdms.TdmsFile(output_file, raw_timestamps=True),
+    )
+
+
 def test_big_endian_format():
     """Test reading a file that encodes data in big endian mode"""
     test_file = tdms.TdmsFile(DATA_DIR + '/big_endian.tdms')

nptdms/timestamp.py

@@ -1,9 +1,11 @@
+import struct
 from datetime import datetime, timedelta
 import numpy as np
 
-
 EPOCH = np.datetime64('1904-01-01 00:00:00', 's')
 
+_struct_pack = struct.pack
+
 
 class TdmsTimestamp(object):
     """ A Timestamp from a TDMS file
@@ -14,6 +16,7 @@ class TdmsTimestamp(object):
         :ivar ~.seconds: Seconds since the epoch as a signed integer
         :ivar ~.second_fractions: A positive number of 2^-64 fractions of a second
     """
+    enum_value = 0x44
 
     def __init__(self, seconds, second_fractions):
         self.seconds = seconds
@@ -27,6 +30,9 @@ def __str__(self):
         fraction_string = "{0:.6f}".format(self.second_fractions * 2.0 ** -64).split('.')[1]
         return "{0}.{1}".format(dt, fraction_string)
 
+    def __eq__(self, other):
+        return self.seconds == other.seconds and self.second_fractions == other.second_fractions
+
     def as_datetime64(self, resolution='us'):
         """ Convert this timestamp to a numpy datetime64 object
 
@@ -49,6 +55,10 @@ def as_datetime(self):
         microseconds = (self.second_fractions / fractions_per_us)
         return datetime(1904, 1, 1, 0, 0, 0) + timedelta(seconds=self.seconds) + timedelta(microseconds=microseconds)
 
+    @property
+    def bytes(self):
+        return _struct_pack('<Qq', self.second_fractions, self.seconds)
+
 
 class TimestampArray(np.ndarray):
     """ A numpy array of TDMS timestamps

nptdms/writer.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 from nptdms.common import toc_properties, ObjectPath
+from nptdms.timestamp import TdmsTimestamp
 from nptdms.types import *
 from nptdms import TdmsFile
 
@@ -38,7 +39,7 @@ def defragment(cls, source, destination, version=4712, index_file=False):
             If ``destination`` is a readable object ``index_file`` can either be a redable object or ``False``
             to store a ``.tdms_index`` file inside of the submitted object or not.
         """
-        file = TdmsFile(source)
+        file = TdmsFile(source, raw_timestamps=True)
         with cls(destination, version=version, index_file=index_file) as new_file:
             new_file.write_segment([RootObject(file.properties)])
             for group in file.groups():
@@ -375,6 +376,8 @@ def _to_tdms_value(value):
         return TimeStamp(value)
     if isinstance(value, np.datetime64):
         return TimeStamp(value)
+    if isinstance(value, TdmsTimestamp):
+        return value
     if isinstance(value, str):
         return String(value)
     if isinstance(value, bytes):