support generics

Author: khvn26

flagsmith/flagsmith.py

@@ -23,6 +23,7 @@
 from flagsmith.types import (
     ApplicationMetadata,
     JsonType,
+    SDKEvaluationContext,
     StreamEvent,
     TraitMapping,
 )
@@ -107,7 +108,7 @@ def __init__(
         self.default_flag_handler = default_flag_handler
         self.enable_realtime_updates = enable_realtime_updates
         self._analytics_processor: typing.Optional[AnalyticsProcessor] = None
-        self._evaluation_context: typing.Optional[engine.EvaluationContext] = None
+        self._evaluation_context: typing.Optional[SDKEvaluationContext] = None
         self._environment_updated_at: typing.Optional[datetime] = None
 
         # argument validation

flagsmith/mappers.py

@@ -6,7 +6,6 @@
 
 import sseclient
 from flag_engine.context.types import (
-    EvaluationContext,
     FeatureContext,
     SegmentContext,
     SegmentRule,
@@ -15,7 +14,12 @@
 from flag_engine.segments.types import ContextValue
 
 from flagsmith.models import Segment
-from flagsmith.types import SegmentMetadata, StreamEvent, TraitConfig
+from flagsmith.types import (
+    SDKEvaluationContext,
+    SegmentMetadata,
+    StreamEvent,
+    TraitConfig,
+)
 
 OverrideKey = typing.Tuple[
     str,
@@ -27,21 +31,20 @@
 
 
 def map_segment_results_to_identity_segments(
-    segment_results: list[SegmentResult],
+    segment_results: list[SegmentResult[SegmentMetadata]],
 ) -> list[Segment]:
     identity_segments: list[Segment] = []
     for segment_result in segment_results:
-        if raw_metadata := segment_result.get("metadata"):
-            metadata = typing.cast(SegmentMetadata, raw_metadata)
-            if metadata.get("source") == "api" and (
-                (flagsmith_id := metadata.get("flagsmith_id")) is not None
-            ):
-                identity_segments.append(
-                    Segment(
-                        id=flagsmith_id,
-                        name=segment_result["name"],
-                    )
+        metadata = segment_result["metadata"]
+        if metadata.get("source") == "api" and (
+            (flagsmith_id := metadata.get("flagsmith_id")) is not None
+        ):
+            identity_segments.append(
+                Segment(
+                    id=flagsmith_id,
+                    name=segment_result["name"],
                 )
+            )
     return identity_segments
 
 
@@ -66,7 +69,7 @@ def map_environment_document_to_environment_updated_at(
 
 
 def map_context_and_identity_data_to_context(
-    context: EvaluationContext,
+    context: SDKEvaluationContext,
     identifier: str,
     traits: typing.Optional[
         typing.Mapping[
@@ -77,7 +80,7 @@ def map_context_and_identity_data_to_context(
             ],
         ]
     ],
-) -> EvaluationContext:
+) -> SDKEvaluationContext:
     return {
         **context,
         "identity": {
@@ -97,7 +100,7 @@ def map_context_and_identity_data_to_context(
 
 def map_environment_document_to_context(
     environment_document: dict[str, typing.Any],
-) -> EvaluationContext:
+) -> SDKEvaluationContext:
     return {
         "environment": {
             "key": environment_document["api_key"],
@@ -122,11 +125,9 @@ def map_environment_document_to_context(
                             segment.get("feature_states") or []
                         )
                     ),
-                    "metadata": dict(
-                        SegmentMetadata(
-                            flagsmith_id=segment_id,
-                            source="api",
-                        )
+                    "metadata": SegmentMetadata(
+                        flagsmith_id=segment_id,
+                        source="api",
                     ),
                 }
                 for segment in environment_document["project"]["segments"]
@@ -140,7 +141,7 @@ def map_environment_document_to_context(
 
 def _map_identity_overrides_to_segments(
     identity_overrides: list[dict[str, typing.Any]],
-) -> dict[str, SegmentContext]:
+) -> dict[str, SegmentContext[SegmentMetadata]]:
     features_to_identifiers: typing.Dict[
         OverridesKey,
         typing.List[str],
@@ -164,12 +165,11 @@ def _map_identity_overrides_to_segments(
             )
         )
         features_to_identifiers[overrides_key].append(identity_override["identifier"])
-    segment_contexts: typing.Dict[str, SegmentContext] = {}
+    segment_contexts: typing.Dict[str, SegmentContext[SegmentMetadata]] = {}
     for overrides_key, identifiers in features_to_identifiers.items():
         # Create a segment context for each unique set of overrides
         # Generate a unique key to avoid collisions
         segment_key = str(hash(overrides_key))
-        segment_metadata = SegmentMetadata(source="identity_overrides")
         segment_contexts[segment_key] = SegmentContext(
             key="",  # Identity override segments never use % Split operator
             name="identity_overrides",
@@ -196,7 +196,7 @@ def _map_identity_overrides_to_segments(
                 }
                 for feature_key, feature_name, feature_enabled, feature_value in overrides_key
             ],
-            metadata=dict(segment_metadata),
+            metadata=SegmentMetadata(source="identity_overrides"),
         )
     return segment_contexts
 

flagsmith/models.py

@@ -3,10 +3,11 @@
 import typing
 from dataclasses import dataclass, field
 
-from flag_engine.result.types import EvaluationResult, FlagResult
+from flag_engine.result.types import FlagResult
 
 from flagsmith.analytics import AnalyticsProcessor
 from flagsmith.exceptions import FlagsmithFeatureDoesNotExistError
+from flagsmith.types import SDKEvaluationResult
 
 
 @dataclass
@@ -57,7 +58,7 @@ class Flags:
     @classmethod
     def from_evaluation_result(
         cls,
-        evaluation_result: EvaluationResult,
+        evaluation_result: SDKEvaluationResult,
         analytics_processor: typing.Optional[AnalyticsProcessor],
         default_flag_handler: typing.Optional[typing.Callable[[str], DefaultFlag]],
     ) -> Flags:

flagsmith/offline_handlers.py

@@ -2,13 +2,12 @@
 from pathlib import Path
 from typing import Protocol
 
-from flag_engine.context.types import EvaluationContext
-
 from flagsmith.mappers import map_environment_document_to_context
+from flagsmith.types import SDKEvaluationContext
 
 
 class OfflineHandler(Protocol):
-    def get_evaluation_context(self) -> EvaluationContext: ...
+    def get_evaluation_context(self) -> SDKEvaluationContext: ...
 
 
 class EvaluationContextLocalFileHandler:
@@ -21,11 +20,11 @@ class EvaluationContextLocalFileHandler:
     """
 
     def __init__(self, file_path: str) -> None:
-        self.evaluation_context: EvaluationContext = json.loads(
+        self.evaluation_context: SDKEvaluationContext = json.loads(
             Path(file_path).read_text(),
         )
 
-    def get_evaluation_context(self) -> EvaluationContext:
+    def get_evaluation_context(self) -> SDKEvaluationContext:
         return self.evaluation_context
 
 
@@ -39,15 +38,15 @@ class EnvironmentDocumentLocalFileHandler:
     """
 
     def __init__(self, file_path: str) -> None:
-        self.evaluation_context: EvaluationContext = (
+        self.evaluation_context: SDKEvaluationContext = (
             map_environment_document_to_context(
                 json.loads(
                     Path(file_path).read_text(),
                 ),
             )
         )
 
-    def get_evaluation_context(self) -> EvaluationContext:
+    def get_evaluation_context(self) -> SDKEvaluationContext:
         return self.evaluation_context
 
 

flagsmith/types.py

@@ -1,7 +1,9 @@
 import typing
 from datetime import datetime
 
+from flag_engine.context.types import EvaluationContext
 from flag_engine.engine import ContextValue
+from flag_engine.result.types import EvaluationResult
 from typing_extensions import NotRequired, TypeAlias
 
 _JsonScalarType: TypeAlias = typing.Union[
@@ -40,3 +42,7 @@ class SegmentMetadata(typing.TypedDict):
     """The ID of the segment used in Flagsmith API."""
     source: NotRequired[typing.Literal["api", "identity_overrides"]]
     """The source of the segment, e.g. 'api', 'identity_overrides'."""
+
+
+SDKEvaluationContext = EvaluationContext[SegmentMetadata]
+SDKEvaluationResult = EvaluationResult[SegmentMetadata]

poetry.lock

@@ -13,7 +13,7 @@ packages = [{ include = "flagsmith" }]
 python = ">=3.9,<4"
 requests = "^2.32.3"
 requests-futures = "^1.0.1"
-flagsmith-flag-engine = "^8.0.0"
+flagsmith-flag-engine = { git = "https://github.com/Flagsmith/flagsmith-engine.git", branch = "feat/generic-metadata" }
 sseclient-py = "^1.8.0"
 
 [tool.poetry.group.dev]

pyproject.toml

@@ -7,13 +7,13 @@
 
 import pytest
 import responses
-from flag_engine.engine import EvaluationContext
 from pyfakefs.fake_filesystem import FakeFilesystem
 from pytest_mock import MockerFixture
 
 from flagsmith import Flagsmith
 from flagsmith.analytics import AnalyticsProcessor
 from flagsmith.mappers import map_environment_document_to_context
+from flagsmith.types import SDKEvaluationContext
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
 
@@ -74,7 +74,7 @@ def local_eval_flagsmith(
 
 
 @pytest.fixture()
-def evaluation_context(environment_json: str) -> EvaluationContext:
+def evaluation_context(environment_json: str) -> SDKEvaluationContext:
     return map_environment_document_to_context(json.loads(environment_json))
 
 

tests/conftest.py

@@ -5,7 +5,6 @@
 import pytest
 import requests
 import responses
-from flag_engine.engine import EvaluationContext
 from pytest_mock import MockerFixture
 from responses import matchers
 
@@ -16,6 +15,7 @@
 )
 from flagsmith.models import DefaultFlag, Flags
 from flagsmith.offline_handlers import OfflineHandler
+from flagsmith.types import SDKEvaluationContext
 
 
 def test_flagsmith_starts_polling_manager_on_init_if_enabled(
@@ -39,7 +39,7 @@ def test_flagsmith_starts_polling_manager_on_init_if_enabled(
 def test_update_environment_sets_environment(
     flagsmith: Flagsmith,
     environment_json: str,
-    evaluation_context: EvaluationContext,
+    evaluation_context: SDKEvaluationContext,
 ) -> None:
     # Given
     responses.add(method="GET", url=flagsmith.environment_url, body=environment_json)
@@ -76,7 +76,7 @@ def test_get_environment_flags_calls_api_when_no_local_environment(
 @responses.activate()
 def test_get_environment_flags_uses_local_environment_when_available(
     flagsmith: Flagsmith,
-    evaluation_context: EvaluationContext,
+    evaluation_context: SDKEvaluationContext,
 ) -> None:
     # Given
     flagsmith._evaluation_context = evaluation_context
@@ -150,7 +150,7 @@ def test_get_identity_flags_calls_api_when_no_local_environment_with_traits(
 @responses.activate()
 def test_get_identity_flags_uses_local_environment_when_available(
     flagsmith: Flagsmith,
-    evaluation_context: EvaluationContext,
+    evaluation_context: SDKEvaluationContext,
     mocker: MockerFixture,
 ) -> None:
     # Given
@@ -545,10 +545,10 @@ def test_initialise_flagsmith_with_proxies() -> None:
     assert flagsmith.session.proxies == proxies
 
 
-def test_offline_mode(evaluation_context: EvaluationContext) -> None:
+def test_offline_mode(evaluation_context: SDKEvaluationContext) -> None:
     # Given
     class DummyOfflineHandler:
-        def get_evaluation_context(self) -> EvaluationContext:
+        def get_evaluation_context(self) -> SDKEvaluationContext:
             return evaluation_context
 
     # When
@@ -566,7 +566,7 @@ def get_evaluation_context(self) -> EvaluationContext:
 @responses.activate()
 def test_flagsmith_uses_offline_handler_if_set_and_no_api_response(
     mocker: MockerFixture,
-    evaluation_context: EvaluationContext,
+    evaluation_context: SDKEvaluationContext,
 ) -> None:
     # Given
     api_url = "http://some.flagsmith.com/api/v1/"
@@ -599,7 +599,7 @@ def test_flagsmith_uses_offline_handler_if_set_and_no_api_response(
 @responses.activate()
 def test_offline_mode__local_evaluation__correct_fallback(
     mocker: MockerFixture,
-    evaluation_context: EvaluationContext,
+    evaluation_context: SDKEvaluationContext,
     caplog: pytest.LogCaptureFixture,
 ) -> None:
     # Given

tests/test_flagsmith.py

@@ -1,9 +1,10 @@
 import typing
 
 import pytest
-from flag_engine.result.types import EvaluationResult, FlagResult
+from flag_engine.result.types import FlagResult
 
 from flagsmith.models import Flag, Flags
+from flagsmith.types import SDKEvaluationResult
 
 
 def test_flag_from_evaluation_result() -> None:
@@ -92,7 +93,7 @@ def test_flags_from_evaluation_result(
     expected_names: typing.List[str],
 ) -> None:
     # Given
-    evaluation_result: EvaluationResult = {
+    evaluation_result: SDKEvaluationResult = {
         "flags": flags_result,
         "segments": [],
     }