fix: Type the sink-or-store helper for each side instead of erasing it

The shared helper took unbound TypeVars, which accepted any sink beside
any store and returned a union, so it checked nothing. Splitting it into
sink_or_store and async_sink_or_store lets each name its own sink and
store types. main did not need this because the async caller sat in an
unannotated method, whose body mypy skips.

Real types surfaced a latent mismatch in _process_message: FEATURES and
SEGMENTS are VersionedDataKindWithOrdering, and Mapping's key type is
invariant, so the inferred dict did not satisfy init's parameter. The
literal now carries the declared type.

Author: jsonbailey

ldclient/impl/datasource/async_polling.py

@@ -9,7 +9,7 @@
 
 from ldclient.async_config import AsyncConfig
 from ldclient.impl.aio.concurrency import AsyncEvent, AsyncRepeatingTask
-from ldclient.impl.datasource.datasource_common import sink_or_store
+from ldclient.impl.datasource.datasource_common import async_sink_or_store
 from ldclient.impl.retry import (
     FailureKind,
     RetryState,
@@ -73,7 +73,7 @@ async def _fetch_and_store(self) -> None:
         """Makes one poll request and records the outcome on the retry state."""
         try:
             all_data = await self._requester.get_all_data()
-            await sink_or_store(self._data_source_update_sink, self._store).init(all_data)
+            await async_sink_or_store(self._data_source_update_sink, self._store).init(all_data)
 
             if self._data_source_update_sink is not None:
                 self._data_source_update_sink.update_status(DataSourceState.VALID, None)

ldclient/impl/datasource/async_streaming.py

@@ -7,7 +7,7 @@
 import asyncio
 import json
 import time
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Mapping, Optional
 from urllib import parse
 
 from ld_eventsource.actions import Event, Fault, Start
@@ -18,8 +18,8 @@
 from ldclient.impl.datasource.datasource_common import (
     STREAM_ALL_PATH,
     StreamClosedError,
-    parse_path,
-    sink_or_store
+    async_sink_or_store,
+    parse_path
 )
 from ldclient.impl.retry import (
     FailureKind,
@@ -34,7 +34,7 @@
     DataSourceErrorKind,
     DataSourceState
 )
-from ldclient.versioned_data_kind import FEATURES, SEGMENTS
+from ldclient.versioned_data_kind import FEATURES, SEGMENTS, VersionedDataKind
 
 
 class AsyncStreamingUpdateProcessor(AsyncUpdateProcessor):
@@ -196,10 +196,10 @@ def initialized(self):
     # Returns True if we initialized the feature store
     async def _process_message(self, msg: Event) -> bool:
         """Process a single SSE event.  Returns True on a successful ``put``."""
-        target = sink_or_store(self._data_source_update_sink, self._store)
+        target = async_sink_or_store(self._data_source_update_sink, self._store)
         if msg.event == 'put':
             all_data = json.loads(msg.data)
-            init_data = {FEATURES: all_data['data']['flags'], SEGMENTS: all_data['data']['segments']}
+            init_data: Mapping[VersionedDataKind, Mapping[str, dict]] = {FEATURES: all_data['data']['flags'], SEGMENTS: all_data['data']['segments']}
             log.debug("Received put event with %d flags and %d segments", len(init_data[FEATURES]), len(init_data[SEGMENTS]))
             await target.init(init_data)
             return True

ldclient/impl/datasource/datasource_common.py

@@ -5,16 +5,15 @@
 # currently excluded from documentation - see docs/README.md
 
 from collections import namedtuple
-from typing import (
-    Mapping,
-    Optional,
-    Protocol,
-    TypeVar,
-    Union,
-    runtime_checkable
-)
+from typing import Mapping, Optional, Protocol, Union, runtime_checkable
 
 from ldclient.impl.util import _LD_ENVID_HEADER
+from ldclient.interfaces import (
+    AsyncDataSourceUpdateSink,
+    AsyncFeatureStore,
+    DataSourceUpdateSink,
+    FeatureStore
+)
 from ldclient.versioned_data_kind import FEATURES, SEGMENTS
 
 STREAM_ALL_PATH = '/all'
@@ -40,11 +39,7 @@ def __init__(self):
         super().__init__("the server closed the stream connection")
 
 
-_Sink = TypeVar('_Sink')
-_Store = TypeVar('_Store')
-
-
-def sink_or_store(sink: Optional[_Sink], store: _Store) -> Union[_Sink, _Store]:
+def sink_or_store(sink: Optional[DataSourceUpdateSink], store: FeatureStore) -> Union[DataSourceUpdateSink, FeatureStore]:
     """
     The original implementation of the data sources relied on the feature store
     directly, which we are trying to move away from. Customers who might have
@@ -62,6 +57,18 @@ def sink_or_store(sink: Optional[_Sink], store: _Store) -> Union[_Sink, _Store]:
     return sink
 
 
+def async_sink_or_store(sink: Optional[AsyncDataSourceUpdateSink], store: AsyncFeatureStore) -> Union[AsyncDataSourceUpdateSink, AsyncFeatureStore]:
+    """
+    The async counterpart of :func:`sink_or_store`, kept separate so each side
+    names its own sink and store types. One shared generic function accepted a
+    sync sink beside an async store without complaint.
+    """
+    if sink is None:
+        return store
+
+    return sink
+
+
 @runtime_checkable
 class EnvironmentIdSink(Protocol):
     """