fix: get_identity_segments returns identity override segments

Author: khvn26

flagsmith/flagsmith.py

@@ -14,6 +14,7 @@
     map_context_and_identity_data_to_context,
     map_environment_document_to_context,
     map_environment_document_to_environment_updated_at,
+    map_segment_results_to_identity_segments,
 )
 from flagsmith.models import DefaultFlag, Flags, Segment
 from flagsmith.offline_handlers import OfflineHandler
@@ -283,10 +284,8 @@ def get_identity_segments(
         evaluation_result = engine.get_evaluation_result(
             context=context,
         )
-        return [
-            Segment(id=int(segment_result["key"]), name=segment_result["name"])
-            for segment_result in evaluation_result["segments"]
-        ]
+
+        return map_segment_results_to_identity_segments(evaluation_result["segments"])
 
     def update_environment(self) -> None:
         try:

flagsmith/mappers.py

@@ -11,9 +11,11 @@
     SegmentContext,
     SegmentRule,
 )
+from flag_engine.result.types import SegmentResult
 from flag_engine.segments.types import ContextValue
 
-from flagsmith.types import StreamEvent, TraitConfig
+from flagsmith.models import Segment
+from flagsmith.types import SegmentMetadata, StreamEvent, TraitConfig
 
 OverrideKey = typing.Tuple[
     str,
@@ -24,6 +26,25 @@
 OverridesKey = typing.Tuple[OverrideKey, ...]
 
 
+def map_segment_results_to_identity_segments(
+    segment_results: list[SegmentResult],
+) -> 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"],
+                    )
+                )
+    return identity_segments
+
+
 def map_sse_event_to_stream_event(event: sseclient.Event) -> StreamEvent:
     event_data = json.loads(event.data)
     return {
@@ -90,7 +111,7 @@ def map_environment_document_to_context(
         },
         "segments": {
             **{
-                (segment_key := str(segment["id"])): {
+                (segment_key := str(segment_id := segment["id"])): {
                     "key": segment_key,
                     "name": segment["name"],
                     "rules": _map_environment_document_rules_to_context_rules(
@@ -101,6 +122,12 @@ def map_environment_document_to_context(
                             segment.get("feature_states") or []
                         )
                     ),
+                    "metadata": dict(
+                        SegmentMetadata(
+                            flagsmith_id=segment_id,
+                            source="api",
+                        )
+                    ),
                 }
                 for segment in environment_document["project"]["segments"]
             },
@@ -142,6 +169,7 @@ def _map_identity_overrides_to_segments(
         # 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",
@@ -168,6 +196,7 @@ def _map_identity_overrides_to_segments(
                 }
                 for feature_key, feature_name, feature_enabled, feature_value in overrides_key
             ],
+            metadata=dict(segment_metadata),
         )
     return segment_contexts
 

flagsmith/types.py

@@ -33,3 +33,10 @@ class TraitConfig(typing.TypedDict):
 class ApplicationMetadata(typing.TypedDict):
     name: NotRequired[str]
     version: NotRequired[str]
+
+
+class SegmentMetadata(typing.TypedDict):
+    flagsmith_id: NotRequired[int]
+    """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'."""

tests/test_flagsmith.py

@@ -509,6 +509,26 @@ def test_get_identity_segments_with_valid_trait(
     assert segments[0].name == "Test segment"  # obtained from data/environment.json
 
 
+def test_get_identity_segments__identity_overrides__returns_expected(
+    local_eval_flagsmith: Flagsmith,
+) -> None:
+    # Given
+    # the identifier matches the identity override in data/environment.json
+    identifier = "overridden-id"
+    # traits match the "Test segment" segment in data/environment.json
+    traits = {"foo": "bar"}
+
+    # When
+    segments = local_eval_flagsmith.get_identity_segments(identifier, traits)
+
+    # Then
+    # identity override virtual segment is not returned,
+    # only the segment matching the traits
+    assert len(segments) == 1
+    assert segments[0].id == 1
+    assert segments[0].name == "Test segment"
+
+
 def test_local_evaluation_requires_server_key() -> None:
     with pytest.raises(ValueError):
         Flagsmith(environment_key="not-a-server-key", enable_local_evaluation=True)