Enums as strings in xarray (#51)

Author: lukasbindreiter

CHANGELOG.md

@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.60.0] - 2026-08-25
+
+### Changed
+
+- `tilebox-datasets`: Dataset queries now return short enum names as strings (for example, `HH`) instead of numeric
+  values with a name mapping in the xarray variable attributes. Repeated enum fields can now be queried and
+  round-tripped through all supported ingestion inputs.
+
 ## [0.59.0] - 2026-08-18
 
 ### Added
@@ -471,7 +479,8 @@ the first client that does not cache data (since it's already on the local file
 - Released under the [MIT](https://opensource.org/license/mit) license.
 - Released packages: `tilebox-datasets`, `tilebox-workflows`, `tilebox-storage`, `tilebox-grpc`
 
-[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.59.0...HEAD
+[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.60.0...HEAD
+[0.60.0]: https://github.com/tilebox/tilebox-python/compare/v0.59.0...v0.60.0
 [0.59.0]: https://github.com/tilebox/tilebox-python/compare/v0.58.0...v0.59.0 
 [0.58.0]: https://github.com/tilebox/tilebox-python/compare/v0.57.0...v0.58.0 
 [0.57.0]: https://github.com/tilebox/tilebox-python/compare/v0.56.0...v0.57.0

tilebox-datasets/tests/protobuf_conversion/test_protobuf_xarray.py

@@ -1,3 +1,4 @@
+from enum import Enum
 from uuid import UUID
 
 import pandas as pd
@@ -25,8 +26,14 @@
 from tilebox.datasets.datasets.stac.v1.core_pb2 import Provider as ProviderPB2
 from tilebox.datasets.datasets.stac.v1.processing_pb import ProcessingSoftware
 from tilebox.datasets.datasets.stac.v1.processing_pb2 import ProcessingSoftware as ProcessingSoftwarePB2
+from tilebox.datasets.datasets.stac.v1.sar_pb2 import (
+    SAR_POLARIZATION_HH,
+    SAR_POLARIZATION_VV,
+    SARProperties,
+)
 from tilebox.datasets.datasets.stac.v1.storage_pb import Storage
 from tilebox.datasets.datasets.stac.v1.storage_pb2 import Storage as StoragePB2
+from tilebox.datasets.datasets.v1.well_known_types_pb2 import ProcessingLevel
 from tilebox.datasets.protobuf_conversion.field_types import _AssetsDisplay
 from tilebox.datasets.protobuf_conversion.protobuf_xarray import MessageToXarrayConverter
 from tilebox.datasets.protobuf_conversion.to_protobuf import to_messages
@@ -87,7 +94,7 @@ def test_convert_datapoint(datapoint: ExampleDatapoint) -> None:  # noqa: PLR091
     )
 
     assert isinstance(dataset.some_geometry.item(), Polygon | MultiPolygon)
-    assert dataset.some_enum.item() == datapoint.some_enum
+    assert dataset.some_enum.item() == ProcessingLevel.Name(datapoint.some_enum).removeprefix("PROCESSING_LEVEL_")
 
     assert list(dataset.some_repeated_string.to_numpy()) == list(datapoint.some_repeated_string)
     assert_array_equal(dataset.some_repeated_int.to_numpy(), datapoint.some_repeated_int)
@@ -213,6 +220,89 @@ def test_convert_stac_messages_to_protobuf_py() -> None:
     assert "access_profiles" not in html
 
 
+def test_convert_scalar_and_repeated_enums_to_names_and_round_trip() -> None:
+    class Polarization(Enum):
+        HH = "HH"
+        VV = "VV"
+
+    file_descriptor = descriptor_pb2.FileDescriptorProto(
+        name="tests/protobuf_conversion/enum_datapoint.proto",
+        package="tests.protobuf_conversion.enums",
+    )
+    enum_descriptor = file_descriptor.enum_type.add(name="SARPolarization")
+    for name, number in (("SAR_POLARIZATION_UNSPECIFIED", 0), ("SAR_POLARIZATION_HH", 1), ("SAR_POLARIZATION_VV", 2)):
+        enum_descriptor.value.add(name=name, number=number)
+    message_descriptor = file_descriptor.message_type.add(name="EnumDatapoint")
+    message_descriptor.field.add(
+        name="primary_polarization",
+        number=1,
+        label=descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL,
+        type=descriptor_pb2.FieldDescriptorProto.TYPE_ENUM,
+        type_name=".tests.protobuf_conversion.enums.SARPolarization",
+    )
+    message_descriptor.field.add(
+        name="polarizations",
+        number=2,
+        label=descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED,
+        type=descriptor_pb2.FieldDescriptorProto.TYPE_ENUM,
+        type_name=".tests.protobuf_conversion.enums.SARPolarization",
+    )
+    descriptor = Default().AddSerializedFile(file_descriptor.SerializeToString())
+    message_type = GetMessageClass(descriptor.message_types_by_name["EnumDatapoint"])
+    messages = [
+        message_type(primary_polarization=1, polarizations=[0, 1]),
+        message_type(primary_polarization=2, polarizations=[2]),
+        message_type(),
+    ]
+
+    converter = MessageToXarrayConverter()
+    converter.convert_all(messages)
+    dataset = converter.finalize("time")
+
+    assert dataset.primary_polarization.dtype == object
+    assert dataset.primary_polarization[:2].to_numpy().tolist() == ["HH", "VV"]
+    assert pd.isna(dataset.primary_polarization[2].item())
+    assert dataset.polarizations.dtype == object
+    assert dataset.polarizations[0].to_numpy().tolist() == ["UNSPECIFIED", "HH"]
+    assert dataset.polarizations[1, 0].item() == "VV"
+    assert pd.isna(dataset.polarizations[1, 1].item())
+    assert pd.isna(dataset.polarizations[2].to_numpy()).all()
+    assert dataset.polarizations.dims == ("time", "n_polarizations")
+    assert dataset.primary_polarization.attrs == {}
+    assert dataset.polarizations.attrs == {}
+    assert to_messages(dataset, message_type) == messages
+
+    expected = message_type(primary_polarization=1, polarizations=[1, 2, 0])
+    record = {"primary_polarization": "HH", "polarizations": ["HH", Polarization.VV, 0]}
+    assert to_messages([record], message_type) == [expected]
+    assert to_messages({name: [value] for name, value in record.items()}, message_type) == [expected]
+    assert to_messages(pd.DataFrame([record]), message_type) == [expected]
+
+    with pytest.raises(ValueError, match="Record 0: Field 'polarizations': Invalid enum name 'INVALID'"):
+        to_messages([{"polarizations": ["INVALID"]}], message_type)
+
+
+def test_convert_sar_polarizations_with_short_names_in_both_directions() -> None:
+    messages = [
+        SARProperties(polarizations=[SAR_POLARIZATION_HH, SAR_POLARIZATION_VV]),
+        SARProperties(polarizations=[SAR_POLARIZATION_VV]),
+        SARProperties(),
+    ]
+    converter = MessageToXarrayConverter()
+    converter.convert_all(messages)
+
+    dataset = converter.finalize("item")
+
+    assert dataset.polarizations[0].to_numpy().tolist() == ["HH", "VV"]
+    assert dataset.polarizations[1, 0].item() == "VV"
+    assert pd.isna(dataset.polarizations[1, 1].item())
+    assert pd.isna(dataset.polarizations[2].to_numpy()).all()
+
+    other_fields = [field.name for field in SARProperties.DESCRIPTOR.fields if field.name != "polarizations"]
+    assert to_messages(dataset, SARProperties, ignore_fields=other_fields) == messages
+    assert to_messages([{"polarizations": ["HH", "VV"]}], SARProperties) == [messages[0]]
+
+
 @given(lists(example_datapoints(generated_fields=True, missing_fields=True), min_size=5, max_size=30))
 def test_convert_datapoints(datapoints: list[ExampleDatapoint]) -> None:  # noqa: C901, PLR0912
     converter = MessageToXarrayConverter()

tilebox-datasets/tests/test_client.py

@@ -145,8 +145,7 @@ def test_find_datapoint() -> None:
 
         if not skip_data:
             assert datapoint.granule_name.item() == "S2A_MSIL1C_20220713T002201_N0400_R102_T08XNS_20220713T015332.SAFE"
-            processing_level = datapoint.processing_level.item()
-            assert datapoint.processing_level.attrs["names"][processing_level] == "L1C"
+            assert datapoint.processing_level.item() == "L1C"
             assert datapoint.copernicus_id.item() == "65505f82-76dd-5e85-b947-a6c879e07446"
             assert isinstance(datapoint.geometry.item(), Polygon)
         else:

tilebox-datasets/tilebox/datasets/protobuf_conversion/field_types.py

@@ -1,5 +1,7 @@
+import re
 from collections.abc import Sequence
 from datetime import timedelta
+from enum import Enum
 from typing import Any
 from uuid import UUID
 
@@ -120,21 +122,30 @@ def to_proto(self, value: Any) -> bool:
 
 class EnumField(ProtobufFieldType):
     def __init__(self, name_lookup: dict[int, str]) -> None:
-        super().__init__(np.uint8)  # we support up to 256 different enum values for now
+        super().__init__(object)
         self._values_to_name = name_lookup
         self._names_to_value = {name: value for value, name in name_lookup.items()}
 
-    def from_proto(self, value: ProtoFieldValue) -> int:
+    def from_proto(self, value: ProtoFieldValue) -> str:
         if not isinstance(value, int):
             raise TypeError(f"Expected int message but got {type(value)}")
-        return value  # we don't parse the value when loading, to avoid having huge arrays of strings
-
-    def to_proto(self, value: str | int) -> int:
+        try:
+            return self._values_to_name[value]
+        except KeyError as error:
+            raise ValueError(f"Invalid enum value {value}") from error
+
+    def to_proto(self, value: str | int | Enum) -> int:
+        if isinstance(value, Enum):
+            value = value.name
         if isinstance(value, (str, np.str_)):
-            return self._names_to_value[value]
-        if int(value) not in self._values_to_name:
+            try:
+                return self._names_to_value[value]
+            except KeyError as error:
+                raise ValueError(f"Invalid enum name {value!r}") from error
+        integer_value = int(value)
+        if integer_value not in self._values_to_name:
             raise ValueError(f"Invalid enum value {value}")  # during ingestion, we can raise an error here
-        return value
+        return integer_value
 
 
 class TimestampField(ProtobufFieldType):
@@ -360,8 +371,11 @@ def _camel_to_uppercase(name: str) -> str:
     Examples:
         >>> _camel_to_uppercase("ProcessingLevel")
         'PROCESSING_LEVEL'
+        >>> _camel_to_uppercase("SARPolarization")
+        'SAR_POLARIZATION'
     """
-    return "".join(["_" + c.lower() if c.isupper() else c for c in name]).lstrip("_").upper()
+    name = re.sub(r"(.)([A-Z][a-z]+)", r"\1_\2", name)
+    return re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", name).upper()
 
 
 def is_missing(value: Any) -> bool:

tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py

@@ -13,10 +13,8 @@
 from numpy.typing import NDArray
 
 from tilebox.datasets.protobuf_conversion.field_types import (
-    EnumField,
     ProtobufFieldType,
     ProtoFieldValue,
-    enum_mapping_from_field_descriptor,
     infer_field_type,
 )
 
@@ -313,26 +311,6 @@ def _resize(self) -> None:
             self._data = data
 
 
-class _EnumFieldConverter(_SimpleFieldConverter):
-    def __init__(self, field_name: str, enum_names: dict[int, str]) -> None:
-        """
-        A field converter for the enum type.
-
-        Args:
-            field_name: The name of enum field in the protobuf message
-        """
-        super().__init__(field_name, EnumField(enum_names))
-        self._enum_names = enum_names
-
-    def finalize(
-        self, dataset: xr.Dataset, count: int, dimension_names: tuple[str, ...], skip_if_empty: bool = False
-    ) -> str | None:
-        field_name = super().finalize(dataset, count, dimension_names, skip_if_empty)
-        if field_name is not None:
-            dataset[field_name].attrs["names"] = self._enum_names
-        return field_name
-
-
 def _create_field_converters(message: Message, buffer_size: int) -> dict[str, _FieldConverter]:
     """
     Create a dictionary mapping from field names to field converters for the given protobuf message descriptor.
@@ -369,13 +347,6 @@ def _create_field_converter(field: FieldDescriptor) -> _FieldConverter:
     Returns:
         A field converter for the given protobuf field descriptor
     """
-    # special handling for enums:
-    if field.type == FieldDescriptor.TYPE_ENUM:
-        if field.is_repeated:
-            raise NotImplementedError("Repeated enum fields are not supported")
-
-        return _EnumFieldConverter(field.name, enum_mapping_from_field_descriptor(field))
-
     field_type = infer_field_type(field)
     if field.is_repeated:
         return _ArrayFieldConverter(field.name, field_type)