From b84c18d806dd75790872896e2225013c9e51290e Mon Sep 17 00:00:00 2001 From: Wei Lu Date: Mon, 31 Aug 2026 16:21:02 -0700 Subject: [PATCH 1/3] python(feat): support enum_types for imports --- python/CHANGELOG.md | 15 +++ .../_tests/resources/test_data_imports.py | 115 ++++++++++++++++++ .../lib/sift_client/sift_types/data_import.py | 87 ++++++------- 3 files changed, 170 insertions(+), 47 deletions(-) diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 55abf79acc..faaaa38fc2 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -22,6 +22,21 @@ data_import = client.data_import.get("data-import-id") Each `DataImport` carries the import's `status`, `error_message`, and `warning_messages`. `job.get_data_import()` returns it for the `Job` that `import_from_path` hands back. +#### Import enum channels from any file format + +Data columns on all import configs now accept `enum_types`, a mapping of enum name to key. + +```python +CsvDataColumn( + column=2, + name="state", + data_type=ChannelDataType.ENUM, + enum_types={"IDLE": 0, "ARMED": 1}, +) +``` + +Setting `enum_types` on a non-enum data type raises a validation error. Multi-channel Parquet single-channel-per-row imports don't support enums, since channel configs there are derived per row. + ## [v0.20.0] - August 25, 2026 ### What's New diff --git a/python/lib/sift_client/_tests/resources/test_data_imports.py b/python/lib/sift_client/_tests/resources/test_data_imports.py index 0be5d898eb..89dd917761 100644 --- a/python/lib/sift_client/_tests/resources/test_data_imports.py +++ b/python/lib/sift_client/_tests/resources/test_data_imports.py @@ -43,8 +43,10 @@ Hdf5ImportConfig, ParquetDataColumn, ParquetFlatDatasetImportConfig, + ParquetSingleChannelConfig, ParquetSingleChannelPerRowImportConfig, ParquetTimeColumn, + TdmsDataColumn, TdmsImportConfig, TimeFormat, UlogDataColumn, @@ -959,3 +961,116 @@ async def test_missing_job_id_raises(self, mock_client, tmp_path): with patch("sift_client.resources.data_imports.upload_file", return_value={}): with pytest.raises(RuntimeError, match="did not include a job ID"): await api.import_from_path(csv_path, config=config, show_progress=False) + + +ENUMS = {"IDLE": 0, "ARMED": 1} + + +def _enum_dict(channel_config: ChannelConfigProto) -> dict[str, int]: + return {e.name: e.key for e in channel_config.enum_types} + + +class TestEnumTypes: + def test_csv_round_trip(self): + config = CsvImportConfig( + asset_name="a", + time_column=CsvTimeColumn(column=1, format=TimeFormat.ABSOLUTE_RFC3339), + data_columns=[ + CsvDataColumn( + column=2, name="state", data_type=ChannelDataType.ENUM, enum_types=ENUMS + ) + ], + ) + proto = config._to_proto() + assert _enum_dict(proto.data_columns[2]) == ENUMS + assert CsvImportConfig._from_proto(proto)["state"].enum_types == ENUMS + + def test_parquet_flat_dataset_round_trip(self): + config = ParquetFlatDatasetImportConfig( + asset_name="a", + time_column=ParquetTimeColumn(path="ts", format=TimeFormat.ABSOLUTE_UNIX_NANOSECONDS), + data_columns=[ + ParquetDataColumn( + path="state", name="state", data_type=ChannelDataType.ENUM, enum_types=ENUMS + ) + ], + ) + proto = config._to_proto() + assert _enum_dict(proto.flat_dataset.data_columns[0].channel_config) == ENUMS + round_tripped = ParquetFlatDatasetImportConfig._from_proto(proto) + assert round_tripped["state"].enum_types == ENUMS + + def test_parquet_single_channel_per_row_round_trip(self): + config = ParquetSingleChannelPerRowImportConfig( + asset_name="a", + time_column=ParquetTimeColumn(path="ts", format=TimeFormat.ABSOLUTE_UNIX_NANOSECONDS), + single_channel=ParquetSingleChannelConfig( + data_path="value", name="state", data_type=ChannelDataType.ENUM, enum_types=ENUMS + ), + ) + proto = config._to_proto() + assert _enum_dict(proto.single_channel_per_row.single_channel.channel) == ENUMS + round_tripped = ParquetSingleChannelPerRowImportConfig._from_proto(proto) + assert round_tripped.single_channel is not None + assert round_tripped.single_channel.enum_types == ENUMS + + def test_tdms_round_trip(self): + config = TdmsImportConfig( + asset_name="a", + data=[ + TdmsDataColumn( + group_name="g", + channel_name="state", + name="state", + data_type=ChannelDataType.ENUM, + enum_types=ENUMS, + ) + ], + ) + proto = config._to_proto() + assert _enum_dict(proto.data[0].channel_config) == ENUMS + assert TdmsImportConfig._from_proto(proto)["state"].enum_types == ENUMS + + def test_hdf5_to_proto(self): + config = Hdf5ImportConfig( + asset_name="a", + time_format=TimeFormat.ABSOLUTE_UNIX_NANOSECONDS, + data=[ + Hdf5DataColumn( + time_dataset="/time", + value_dataset="/state", + name="state", + data_type=ChannelDataType.ENUM, + enum_types=ENUMS, + ) + ], + ) + proto = config._to_proto() + assert _enum_dict(proto.data[0].channel_config) == ENUMS + + def test_ulog_round_trip(self): + config = UlogImportConfig( + asset_name="a", + data=[ + UlogDataColumn( + message_name="vehicle_status", + field_name="arming_state", + data_type=ChannelDataType.ENUM, + enum_types=ENUMS, + ) + ], + ) + proto = config._to_proto() + assert _enum_dict(proto.data[0].channel_config) == ENUMS + round_tripped = UlogImportConfig._from_proto(proto) + assert round_tripped["vehicle_status_0.arming_state"].enum_types == ENUMS + + def test_enum_types_require_enum_data_type(self): + with pytest.raises(ValueError, match="ENUM"): + CsvDataColumn( + column=2, name="state", data_type=ChannelDataType.DOUBLE, enum_types=ENUMS + ) + + def test_no_enum_types_from_proto_is_none(self, csv_config): + round_tripped = CsvImportConfig._from_proto(csv_config._to_proto()) + assert round_tripped["cpu_util"].enum_types is None diff --git a/python/lib/sift_client/sift_types/data_import.py b/python/lib/sift_client/sift_types/data_import.py index 94292e69bb..cf8225cf3f 100644 --- a/python/lib/sift_client/sift_types/data_import.py +++ b/python/lib/sift_client/sift_types/data_import.py @@ -144,12 +144,41 @@ class DataColumnBase(BaseModel, ABC): data_type: The data type of the channel values. units: Optional units string. description: Optional channel description. + enum_types: Mapping of enum name to key. Only valid when ``data_type`` + is ``ChannelDataType.ENUM``. """ name: str data_type: ChannelDataType units: str = "" description: str = "" + enum_types: dict[str, int] | None = None + + @model_validator(mode="after") + def _check_enum_types(self) -> DataColumnBase: + if self.enum_types and self.data_type != ChannelDataType.ENUM: + raise ValueError( + f"'enum_types' requires data_type ChannelDataType.ENUM " + f"({self.name} is {self.data_type.name})." + ) + return self + + def _channel_config_proto(self) -> ChannelConfigProto: + proto = ChannelConfigProto( + name=self.name, + data_type=self.data_type.value, + units=self.units, + description=self.description, + ) + if self.enum_types: + proto.enum_types.extend( + ChannelEnumTypeProto(name=name, key=key) for name, key in self.enum_types.items() + ) + return proto + + +def _enum_types_from_proto(channel_config: ChannelConfigProto) -> dict[str, int] | None: + return {e.name: e.key for e in channel_config.enum_types} or None class ImportConfigBase(BaseModel, ABC): @@ -229,15 +258,7 @@ def _to_proto(self) -> CsvConfigProto: run_id=self.run_id or "", first_data_row=self.first_data_row, time_column=self.time_column._to_proto(), - data_columns={ - dc.column: ChannelConfigProto( - name=dc.name, - data_type=dc.data_type.value, - units=dc.units, - description=dc.description, - ) - for dc in self.data_columns - }, + data_columns={dc.column: dc._channel_config_proto() for dc in self.data_columns}, ) @classmethod @@ -262,6 +283,7 @@ def _from_proto(cls, proto: CsvConfigProto) -> CsvImportConfig: data_type=ChannelDataType(ch_cfg.data_type), units=ch_cfg.units, description=ch_cfg.description, + enum_types=_enum_types_from_proto(ch_cfg), ) for col_num, ch_cfg in proto.data_columns.items() ] @@ -383,12 +405,7 @@ def _to_proto(self) -> ParquetConfigProto: data_columns=[ ParquetDataColumnProto( path=dc.path, - channel_config=ChannelConfigProto( - name=dc.name, - data_type=dc.data_type.value, - units=dc.units, - description=dc.description, - ), + channel_config=dc._channel_config_proto(), ) for dc in self.data_columns ], @@ -420,6 +437,7 @@ def _from_proto( data_type=ChannelDataType(dc.channel_config.data_type), units=dc.channel_config.units, description=dc.channel_config.description, + enum_types=_enum_types_from_proto(dc.channel_config), ) for dc in fd.data_columns ] @@ -509,12 +527,7 @@ def _to_proto(self) -> ParquetConfigProto: scpr.single_channel.CopyFrom( ParquetSingleChannelPerRowSingleChannelConfigProto( data_path=sc.data_path, - channel=ChannelConfigProto( - name=sc.name, - data_type=sc.data_type.value, - units=sc.units, - description=sc.description, - ), + channel=sc._channel_config_proto(), ) ) elif self.multi_channel is not None: @@ -556,6 +569,7 @@ def _from_proto( data_type=ChannelDataType(sc.channel.data_type), units=sc.channel.units, description=sc.channel.description, + enum_types=_enum_types_from_proto(sc.channel), ) elif scpr.HasField("multi_channel"): mc = scpr.multi_channel @@ -629,7 +643,6 @@ class TdmsDataColumn(DataColumnBase): time_channel_name: str | None = None scaled: bool | None = None complex_component: TdmsComplexComponent | None = None - enum_types: dict[str, int] | None = None class TdmsImportConfig(ImportConfigBase): @@ -674,20 +687,10 @@ def _to_proto(self) -> TDMSConfigProto: if self.relative_start_time is not None: proto.relative_start_time.CopyFrom(to_pb_timestamp(self.relative_start_time)) for d in self.data: - channel_config = ChannelConfigProto( - name=d.name, - data_type=d.data_type.value, - units=d.units, - description=d.description, - ) - if d.enum_types: - channel_config.enum_types.extend( - ChannelEnumTypeProto(name=name, key=key) for name, key in d.enum_types.items() - ) entry = TdmsDataConfigProto( group_name=d.group_name, channel_name=d.channel_name, - channel_config=channel_config, + channel_config=d._channel_config_proto(), ) if d.time_channel_name is not None: entry.time_channel_name = d.time_channel_name @@ -719,7 +722,6 @@ def _from_proto(cls, proto: TDMSConfigProto) -> TdmsImportConfig: complex_component = None if d.complex_component and d.complex_component != TDMS_COMPLEX_COMPONENT_UNSPECIFIED: complex_component = TdmsComplexComponent(d.complex_component) - enum_types = {e.name: e.key for e in ch.enum_types} if ch.enum_types else None data.append( TdmsDataColumn( group_name=d.group_name, @@ -733,7 +735,7 @@ def _from_proto(cls, proto: TDMSConfigProto) -> TdmsImportConfig: else None, scaled=d.scaled if d.HasField("scaled") else None, complex_component=complex_component, - enum_types=enum_types, + enum_types=_enum_types_from_proto(ch), ) ) @@ -838,12 +840,7 @@ def _to_proto(self) -> Hdf5ConfigProto: time_index=d.time_index, value_dataset=d.value_dataset, value_index=d.value_index, - channel_config=ChannelConfigProto( - name=d.name, - data_type=d.data_type.value, - units=d.units, - description=d.description, - ), + channel_config=d._channel_config_proto(), time_field=d.time_field, value_field=d.value_field, ) @@ -967,12 +964,7 @@ def _to_proto(self) -> UlogConfigProto: message_name=dc.message_name, instance=dc.instance, field_name=dc.field_name, - channel_config=ChannelConfigProto( - name=dc.name, - data_type=dc.data_type.value, - units=dc.units, - description=dc.description, - ), + channel_config=dc._channel_config_proto(), ) ) return proto @@ -999,6 +991,7 @@ def _from_proto(cls, proto: UlogConfigProto) -> UlogImportConfig: data_type=ChannelDataType(d.channel_config.data_type), units=d.channel_config.units, description=d.channel_config.description, + enum_types=_enum_types_from_proto(d.channel_config), ) for d in proto.data ] From 2d1b9b12167e27706fff0af0c3b24baafda9ce00 Mon Sep 17 00:00:00 2001 From: Wei Lu Date: Mon, 31 Aug 2026 16:55:19 -0700 Subject: [PATCH 2/3] raise on duplicate keys --- .../_tests/resources/test_data_imports.py | 9 +++++++++ python/lib/sift_client/sift_types/data_import.py | 13 ++++++++----- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/python/lib/sift_client/_tests/resources/test_data_imports.py b/python/lib/sift_client/_tests/resources/test_data_imports.py index 89dd917761..d2f8666073 100644 --- a/python/lib/sift_client/_tests/resources/test_data_imports.py +++ b/python/lib/sift_client/_tests/resources/test_data_imports.py @@ -1071,6 +1071,15 @@ def test_enum_types_require_enum_data_type(self): column=2, name="state", data_type=ChannelDataType.DOUBLE, enum_types=ENUMS ) + def test_enum_types_reject_duplicate_keys(self): + with pytest.raises(ValueError, match="[Dd]uplicate"): + CsvDataColumn( + column=2, + name="state", + data_type=ChannelDataType.ENUM, + enum_types={"IDLE": 0, "ARMED": 0}, + ) + def test_no_enum_types_from_proto_is_none(self, csv_config): round_tripped = CsvImportConfig._from_proto(csv_config._to_proto()) assert round_tripped["cpu_util"].enum_types is None diff --git a/python/lib/sift_client/sift_types/data_import.py b/python/lib/sift_client/sift_types/data_import.py index cf8225cf3f..a71e77e268 100644 --- a/python/lib/sift_client/sift_types/data_import.py +++ b/python/lib/sift_client/sift_types/data_import.py @@ -156,11 +156,14 @@ class DataColumnBase(BaseModel, ABC): @model_validator(mode="after") def _check_enum_types(self) -> DataColumnBase: - if self.enum_types and self.data_type != ChannelDataType.ENUM: - raise ValueError( - f"'enum_types' requires data_type ChannelDataType.ENUM " - f"({self.name} is {self.data_type.name})." - ) + if self.enum_types: + if self.data_type != ChannelDataType.ENUM: + raise ValueError( + f"'enum_types' requires data_type ChannelDataType.ENUM " + f"({self.name} is {self.data_type.name})." + ) + if len(set(self.enum_types.values())) != len(self.enum_types): + raise ValueError(f"'enum_types' contains duplicate keys ({self.name}).") return self def _channel_config_proto(self) -> ChannelConfigProto: From 9722753a4a3990ef1a606f114516499298bcdb0d Mon Sep 17 00:00:00 2001 From: Wei Lu Date: Tue, 1 Sep 2026 11:20:15 -0700 Subject: [PATCH 3/3] check mapping is included for enums --- .../lib/sift_client/_tests/resources/test_data_imports.py | 4 ++++ python/lib/sift_client/sift_types/data_import.py | 6 +++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/python/lib/sift_client/_tests/resources/test_data_imports.py b/python/lib/sift_client/_tests/resources/test_data_imports.py index d2f8666073..893b504dbb 100644 --- a/python/lib/sift_client/_tests/resources/test_data_imports.py +++ b/python/lib/sift_client/_tests/resources/test_data_imports.py @@ -1071,6 +1071,10 @@ def test_enum_types_require_enum_data_type(self): column=2, name="state", data_type=ChannelDataType.DOUBLE, enum_types=ENUMS ) + def test_enum_data_type_requires_enum_types(self): + with pytest.raises(ValueError, match="requires 'enum_types'"): + CsvDataColumn(column=2, name="state", data_type=ChannelDataType.ENUM) + def test_enum_types_reject_duplicate_keys(self): with pytest.raises(ValueError, match="[Dd]uplicate"): CsvDataColumn( diff --git a/python/lib/sift_client/sift_types/data_import.py b/python/lib/sift_client/sift_types/data_import.py index a71e77e268..1966109dfa 100644 --- a/python/lib/sift_client/sift_types/data_import.py +++ b/python/lib/sift_client/sift_types/data_import.py @@ -156,12 +156,16 @@ class DataColumnBase(BaseModel, ABC): @model_validator(mode="after") def _check_enum_types(self) -> DataColumnBase: - if self.enum_types: + if self.data_type == ChannelDataType.ENUM or self.enum_types: if self.data_type != ChannelDataType.ENUM: raise ValueError( f"'enum_types' requires data_type ChannelDataType.ENUM " f"({self.name} is {self.data_type.name})." ) + if not self.enum_types: + raise ValueError( + f"data_type ChannelDataType.ENUM requires 'enum_types' ({self.name})." + ) if len(set(self.enum_types.values())) != len(self.enum_types): raise ValueError(f"'enum_types' contains duplicate keys ({self.name}).") return self