- fix tests

- Properly handle optional dependency of sift-stream

Author: alexluck-sift

python/lib/sift_client/_internal/low_level_wrappers/ingestion.py

@@ -18,17 +18,14 @@
 from queue import Queue
 from typing import TYPE_CHECKING, Any, cast
 
-import sift_stream_bindings
 from sift.ingestion_configs.v2.ingestion_configs_pb2 import (
     GetIngestionConfigRequest,
     ListIngestionConfigFlowsResponse,
     ListIngestionConfigsRequest,
     ListIngestionConfigsResponse,
 )
-from sift.ingestion_configs.v2.ingestion_configs_pb2_grpc import IngestionConfigServiceStub
-from sift_stream_bindings import (
-    IngestionConfigFormPy,
-    IngestWithConfigDataStreamRequestPy,
+from sift.ingestion_configs.v2.ingestion_configs_pb2_grpc import (
+    IngestionConfigServiceStub,
 )
 
 from sift_client._internal.low_level_wrappers.base import (
@@ -44,6 +41,12 @@
 if TYPE_CHECKING:
     from datetime import datetime
 
+    from sift_stream_bindings import (
+        IngestionConfigFormPy,
+        IngestWithConfigDataStreamRequestPy,
+        SiftStreamBuilderPy,
+    )
+
 
 class IngestionThread(threading.Thread):
     """Manages ingestion for a single ingestion config."""
@@ -54,7 +57,7 @@ class IngestionThread(threading.Thread):
 
     def __init__(
         self,
-        sift_stream_builder: sift_stream_bindings.SiftStreamBuilderPy,
+        sift_stream_builder: SiftStreamBuilderPy,
         data_queue: Queue,
         ingestion_config: IngestionConfigFormPy,
         no_data_timeout: int = 1,
@@ -154,7 +157,7 @@ class IngestionLowLevelClient(LowLevelClientBase, WithGrpcClient):
 
     CacheEntry = namedtuple("CacheEntry", ["data_queue", "ingestion_config", "thread"])
 
-    sift_stream_builder: sift_stream_bindings.SiftStreamBuilderPy
+    sift_stream_builder: SiftStreamBuilderPy
     stream_cache: dict[str, CacheEntry]
 
     def __init__(self, grpc_client: GrpcClient):
@@ -163,21 +166,25 @@ def __init__(self, grpc_client: GrpcClient):
         Args:
             grpc_client: The gRPC client to use for making API calls.
         """
+        from sift_stream_bindings import (
+            RecoveryStrategyPy,
+            RetryPolicyPy,
+            SiftStreamBuilderPy,
+        )
+
         super().__init__(grpc_client=grpc_client)
         # Rust GRPC client expects URI to have http(s):// prefix.
         uri = grpc_client._config.uri
         if not uri.startswith("http"):
             uri = f"https://{uri}" if grpc_client._config.use_ssl else f"http://{uri}"
-        self.sift_stream_builder = sift_stream_bindings.SiftStreamBuilderPy(
+        self.sift_stream_builder = SiftStreamBuilderPy(
             uri=uri,
             apikey=grpc_client._config.api_key,
         )
         self.sift_stream_builder.enable_tls = grpc_client._config.use_ssl
         # FD-177: Expose configuration for recovery strategy.
-        self.sift_stream_builder.recovery_strategy = (
-            sift_stream_bindings.RecoveryStrategyPy.retry_only(
-                sift_stream_bindings.RetryPolicyPy.default()
-            )
+        self.sift_stream_builder.recovery_strategy = RecoveryStrategyPy.retry_only(
+            RetryPolicyPy.default()
         )
         self.stream_cache = {}
 
@@ -229,7 +236,9 @@ async def get_ingestion_config_id_from_client_key(self, client_key: str) -> str
         return ingestion_configs[0].id_
 
     def _new_ingestion_thread(
-        self, ingestion_config_id: str, ingestion_config: IngestionConfigFormPy | None = None
+        self,
+        ingestion_config_id: str,
+        ingestion_config: IngestionConfigFormPy | None = None,
     ):
         """Start a new ingestion thread.
         This allows ingestion to happen in the background regardless of if the user is using the sync or async client

python/lib/sift_client/_tests/sift_types/test_base.py

@@ -41,7 +41,7 @@ class NestedCreateModel(ModelCreate[CreateCalculatedChannelRequest]):
     expression: str | None = None
     all_assets: bool | None = None
 
-    _to_proto_helpers: ClassVar = {
+    _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = {
         "expression": MappingHelper(
             proto_attr_path="calculated_channel_configuration.query_configuration.sel.expression",
             update_field="query_configuration",
@@ -293,7 +293,7 @@ def test_validation_error_on_invalid_helper_field(self):
             class InvalidModel(ModelCreate[CreateCalculatedChannelRequest]):
                 name: str
 
-                _to_proto_helpers: ClassVar = {
+                _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = {
                     "nonexistent_field": MappingHelper(proto_attr_path="some.path"),
                 }
 

python/lib/sift_client/sift_types/_base.py

@@ -2,7 +2,14 @@
 
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, TypeVar
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    ClassVar,
+    Generic,
+    TypeVar,
+)
 
 from google.protobuf import field_mask_pb2, message
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator
@@ -84,17 +91,31 @@ class ModelCreateUpdateBase(BaseModel, ABC):
     """Base class for Pydantic models that generate proto messages."""
 
     model_config = ConfigDict(frozen=False)
-    _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = PrivateAttr(default={})
+    _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = {}
 
     def __init__(self, **data: Any):
         super().__init__(**data)
+
+    @model_validator(mode="after")
+    def _check_mapping_helpers(self):
         if self._to_proto_helpers:
             data = self.model_dump()
             for expected_field in self._to_proto_helpers.keys():
                 if expected_field not in data:
                     raise ValueError(
                         f"MappingHelper created for {expected_field} but {self.__class__.__name__} has no matching variable names."
                     )
+        return self
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        required_annotation = "ClassVar[dict[str, MappingHelper]]"
+        annotation = cls.__annotations__.get("_to_proto_helpers")
+        # Check for correct annotation otherwise pydantic will not populate this properly
+        if annotation and annotation != required_annotation:
+            raise TypeError(
+                f"{cls.__name__} must define _to_proto_helpers type as: {required_annotation}"
+            )
 
     def _build_proto_and_paths(
         self, proto_msg, data, prefix="", already_setting_path_override=False
@@ -123,13 +144,13 @@ def _build_proto_and_paths(
                 if mapping_helper.update_field:
                     paths.append(mapping_helper.update_field)
             elif isinstance(value, dict):
-                if field_name in self._to_proto_helpers:
-                    assert self._to_proto_helpers[field_name].converter, (
+                if field_name in self.__class__._to_proto_helpers:
+                    assert self.__class__._to_proto_helpers[field_name].converter, (
                         f"Expecting to run a coverter given a helper was defined for: {field_name}"
                     )
                     sub_paths = self._build_proto_and_paths(
                         proto_msg,
-                        {field_name: self._to_proto_helpers[field_name].converter(value)},  # type: ignore[misc]
+                        {field_name: self.__class__._to_proto_helpers[field_name].converter(value)},  # type: ignore[misc]
                         "",
                         already_setting_path_override=True,
                     )
@@ -151,13 +172,13 @@ def _build_proto_and_paths(
                 try:
                     repeated_field.extend(value)  # Add all new values
                 except TypeError as e:
-                    if field_name in self._to_proto_helpers:
-                        assert self._to_proto_helpers[field_name].converter, (
+                    if field_name in self.__class__._to_proto_helpers:
+                        assert self.__class__._to_proto_helpers[field_name].converter, (
                             f"Expecting to run a coverter given a helper was defined for: {field_name}"
                         )
                         for item in value:
                             repeated_field.append(
-                                self._to_proto_helpers[field_name].converter(**item)  # type: ignore
+                                self.__class__._to_proto_helpers[field_name].converter(**item)  # type: ignore
                             )
                     else:
                         raise e

python/lib/sift_client/sift_types/asset.py

@@ -113,7 +113,7 @@ class AssetUpdate(ModelUpdate[AssetProto]):
     metadata: dict[str, str | float | bool] | None = None
     is_archived: bool | None = None
 
-    _to_proto_helpers: ClassVar = {
+    _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = {
         "metadata": MappingHelper(
             proto_attr_path="metadata",
             update_field="metadata",

python/lib/sift_client/sift_types/calculated_channel.py

@@ -149,7 +149,7 @@ class CalculatedChannelBase(ModelCreateUpdateBase):
 
     metadata: dict[str, str | float | bool] | None = None
 
-    _to_proto_helpers: ClassVar = {
+    _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = {
         "expression": MappingHelper(
             proto_attr_path="calculated_channel_configuration.query_configuration.sel.expression",
             update_field="query_configuration",

python/lib/sift_client/sift_types/ingestion.py

@@ -15,21 +15,20 @@
 from sift.ingestion_configs.v2.ingestion_configs_pb2 import (
     IngestionConfig as IngestionConfigProto,
 )
-from sift_stream_bindings import (
-    ChannelBitFieldElementPy,
-    ChannelConfigPy,
-    ChannelDataTypePy,
-    ChannelEnumTypePy,
-    FlowConfigPy,
-    IngestWithConfigDataChannelValuePy,
-)
 
 from sift_client.sift_types._base import BaseType
 from sift_client.sift_types.channel import ChannelBitFieldElement, ChannelDataType
 
 if TYPE_CHECKING:
     from datetime import datetime
 
+    from sift_stream_bindings import (
+        ChannelConfigPy,
+        ChannelDataTypePy,
+        FlowConfigPy,
+        IngestWithConfigDataChannelValuePy,
+    )
+
     from sift_client.client import SiftClient
     from sift_client.sift_types.channel import Channel
 
@@ -179,6 +178,8 @@ def _to_proto(self) -> FlowConfig:
         )
 
     def _to_rust_config(self) -> FlowConfigPy:
+        from sift_stream_bindings import FlowConfigPy
+
         return FlowConfigPy(
             name=self.name,
             channels=[_channel_to_rust_config(channel) for channel in self.channels],
@@ -218,6 +219,12 @@ def ingest(self, *, timestamp: datetime, channel_values: dict[str, Any]):
 
 # Converter functions.
 def _channel_to_rust_config(channel: ChannelConfig) -> ChannelConfigPy:
+    from sift_stream_bindings import (
+        ChannelBitFieldElementPy,
+        ChannelConfigPy,
+        ChannelEnumTypePy,
+    )
+
     return ChannelConfigPy(
         name=channel.name,
         data_type=_to_rust_type(channel.data_type),
@@ -252,6 +259,8 @@ def _rust_channel_value_from_bitfield(
     Returns:
         A ChannelValuePy object.
     """
+    from sift_stream_bindings import IngestWithConfigDataChannelValuePy
+
     assert channel.bit_field_elements is not None
     # We expect individual ints or bytes to represent full bitfield values.
     if isinstance(value, bytes) or isinstance(value, int):
@@ -276,6 +285,8 @@ def _rust_channel_value_from_bitfield(
 
 
 def _to_rust_value(channel: ChannelConfig, value: Any) -> IngestWithConfigDataChannelValuePy:
+    from sift_stream_bindings import IngestWithConfigDataChannelValuePy
+
     if value is None:
         return IngestWithConfigDataChannelValuePy.empty()
     if channel.data_type == ChannelDataType.ENUM and channel.enum_types is not None:
@@ -314,6 +325,8 @@ def _to_rust_value(channel: ChannelConfig, value: Any) -> IngestWithConfigDataCh
 
 
 def _to_rust_type(data_type: ChannelDataType) -> ChannelDataTypePy:
+    from sift_stream_bindings import ChannelDataTypePy
+
     if data_type == ChannelDataType.DOUBLE:
         return ChannelDataTypePy.Double
     elif data_type == ChannelDataType.FLOAT:

python/lib/sift_client/sift_types/run.py

@@ -126,7 +126,7 @@ class RunBase(ModelCreateUpdateBase):
     tags: list[str] | None = None
     metadata: dict[str, str | float | bool] | None = None
 
-    _to_proto_helpers: ClassVar = {
+    _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = {
         "metadata": MappingHelper(
             proto_attr_path="metadata",
             update_field="metadata",
@@ -158,7 +158,7 @@ class RunCreate(RunBase, ModelCreate[CreateRunRequestProto]):
     stop_time: datetime | None = None
     organization_id: str | None = None
 
-    _to_proto_helpers: ClassVar = {
+    _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = {
         "metadata": MappingHelper(
             proto_attr_path="metadata",
             update_field="metadata",
@@ -180,7 +180,7 @@ class RunUpdate(RunBase, ModelUpdate[RunProto]):
     stop_time: datetime | None = None
     is_archived: bool | None = None
 
-    _to_proto_helpers: ClassVar = {
+    _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = {
         "metadata": MappingHelper(
             proto_attr_path="metadata",
             update_field="metadata",

python/pyproject.toml

@@ -22,7 +22,7 @@ dependencies = [
     "PyYAML~=6.0",
     "pandas~=2.0",
     "protobuf>=4.0",
-    "pydantic~=2.0",
+    "pydantic~=2.10",
     # Support python 3.9+ typing in older versons of python.
     "eval-type-backport~=0.2",
     "pydantic_core~=2.3",