feat: adding ability to specify time column in TDMS groups

nominal/nominal.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from datetime import datetime
 from functools import cache
 from pathlib import Path
@@ -28,6 +29,7 @@
     import pandas as pd
     import polars as pl
 
+logger = logging.getLogger(__name__)
 
 _DEFAULT_BASE_URL = "https://api.gov.nominal.io/api"
 
@@ -74,33 +76,114 @@ def get_user() -> User:
 
 
 def upload_tdms(
-    file: Path | str, name: str | None = None, description: str | None = None, *, wait_until_complete: bool = True
+    file: Path | str, name: str | None = None, description: str | None = None, timestamp_column: str | None = None,
+        timestamp_type: ts._AnyTimestampType | None = None, *, wait_until_complete: bool = True
 ) -> Dataset:
     """Create a dataset in the Nominal platform from a tdms file.
 
-    TDMS channel properties must have both a `wf_increment` and `wf_start_time` property to be included in the dataset.
+    If `name` is None, the dataset is created with the name of the file with a .csv suffix.
 
-    Channels will be named as f"{group_name}.{channel_name}" with spaces replaced with underscores.
+    If 'timestamp_column' is provided, it must be present in every group and the length of all data columns must be
+    equal to (and aligned with) with 'timestamp_column'.
 
-    If `name` is None, the dataset is created with the name of the file with a .csv suffix.
+    If 'timestamp_column' is None, TDMS channel properties must have both a `wf_increment` and `wf_start_time`
+    property to be included in the dataset.
+
+    Channels will be named as f"{group_name}.{channel_name}" with spaces replaced with underscores.
 
     If `wait_until_complete=True` (the default), this function waits until the dataset has completed ingestion before
         returning. If you are uploading many datasets, set `wait_until_complete=False` instead and call
         `wait_until_ingestions_complete()` after uploading all datasets to allow for parallel ingestion.
     """
-    import numpy as np
+
+    if timestamp_column:
+        return _tdms_upload_with_time_column(file, name, description, timestamp_column, timestamp_type,
+                                             wait_until_complete=wait_until_complete)
+    else:
+        return _tdms_upload_with_waveform_props(file, name, description, wait_until_complete=wait_until_complete)
+
+
+def _tdms_upload_with_time_column(file: Path | str, name: str | None = None, description: str | None = None,
+                                  timestamp_column: str | None = None,
+                                  timestamp_type: ts._AnyTimestampType | None = None, *,
+                                  wait_until_complete: bool = True):
+    from nptdms import TdmsChannel, TdmsFile, TdmsGroup
     import pandas as pd
+
+    path = Path(file)
+    with TdmsFile.open(path) as tdms_file:
+        channels_to_export: dict[str, TdmsChannel] = {}
+        group: TdmsGroup
+        df = None
+
+        for group in tdms_file.groups():
+            time_channel = None
+            channel: TdmsChannel
+
+            # pull out timestamp column if expected
+            for channel in group.channels():
+                if channel.name == timestamp_column:
+                    time_channel = channel
+            if not time_channel:
+                logger.info(
+                    "Skipping channel group \"%s\" because expected timestamp_column \"%s\" does not exist",
+                    channel.group_name, timestamp_column)
+                continue
+
+            # select channels for export
+            for channel in group.channels():
+                # do not export timestamp column
+                if channel.name == timestamp_column:
+                    continue
+
+                # skip if unexpected column length
+                if len(channel) != len(time_channel):
+                    logger.info("Skipping channel \"%s\" because length does not match \"%s\"",
+                                channel.name, timestamp_column)
+                    continue
+                channel_name = f"{channel.group_name.replace(' ', '_')}.{channel.name.replace(' ', '_')}"
+                channels_to_export[channel_name] = channel
+
+            group_df = pd.DataFrame.from_dict(
+                {
+                    channel_name: pd.Series(
+                        data=channel.read_data(), index=time_channel.read_data()
+                    )
+                    for channel_name, channel in channels_to_export.items()
+                })
+            df = group_df if not df else pd.merge(left=df, right=group_df, left_index=True, right_index=True,
+                                                  how='outer')
+
+        # format for nominal upload
+        df.index = df.index.set_names(timestamp_column, level=None)
+        df = df.reset_index()
+
+        return upload_pandas(
+            df=df,
+            name=name if name is not None else path.with_suffix(".csv").name,
+            description=description,
+            timestamp_column=timestamp_column,
+            timestamp_type=timestamp_type,
+            wait_until_complete=wait_until_complete,
+        )
+
+
+def _tdms_upload_with_waveform_props(file: Path | str, name: str | None = None, description: str | None = None, *,
+                                     wait_until_complete: bool = True):
     from nptdms import TdmsChannel, TdmsFile, TdmsGroup
+    import numpy as np
+    import pandas as pd
 
     path = Path(file)
     with TdmsFile.open(path) as tdms_file:
-        # identify channels to extract
         channels_to_export: dict[str, TdmsChannel] = {}
         group: TdmsGroup
+
+        # select channels for export
         for group in tdms_file.groups():
             channel: TdmsChannel
             for channel in group.channels():
-                # some channels may not have the required properties to construct a time track
+                # skip channel if it does not have the required waveform properties to construct a time track
                 if ("wf_increment" in channel.properties) and ("wf_start_time" in channel.properties):
                     channel_name = f"{channel.group_name.replace(' ', '_')}.{channel.name.replace(' ', '_')}"
                     channels_to_export[channel_name] = channel