Merge branch 'perf/local-eval-regression' into feat/flags-context-values

# Conflicts:
#	flagsmith/local_evaluation.py
#	tests/test_flagsmith.py

Author: khvn26

flagsmith/flagsmith.py

@@ -30,7 +30,6 @@
 from flagsmith.types import (
     ApplicationMetadata,
     JsonType,
-    SDKEvaluationContext,
     StreamEvent,
     TraitMapping,
 )
@@ -129,7 +128,6 @@ def __init__(
         self.enable_realtime_updates = enable_realtime_updates
         self._analytics_processor: typing.Optional[AnalyticsProcessor] = None
         self._event_processor: typing.Optional[EventProcessor] = None
-        self.__evaluation_context: typing.Optional[SDKEvaluationContext] = None
         self._local_evaluation: typing.Optional[LocalEvaluation] = None
         self._environment_updated_at: typing.Optional[datetime] = None
 
@@ -152,8 +150,10 @@ def __init__(
             )
 
         if self.offline_handler:
-            self._evaluation_context = map_environment_document_to_context(
-                self.offline_handler.get_environment()
+            self._local_evaluation = LocalEvaluation(
+                map_environment_document_to_context(
+                    self.offline_handler.get_environment()
+                )
             )
 
         if not self.offline_mode:
@@ -242,12 +242,13 @@ def _initialise_local_evaluation(self) -> None:
         # method calls, update the environment manually.
         self.update_environment()
         if self.enable_realtime_updates:
-            if not self._evaluation_context:
+            if not self._local_evaluation:
                 raise ValueError("Unable to get environment from API key")
 
+            environment_key = self._local_evaluation.context["environment"]["key"]
             stream_url = urljoin(
                 self.realtime_api_url,
-                f"sse/environments/{self._evaluation_context['environment']['key']}/stream",
+                f"sse/environments/{environment_key}/stream",
             )
 
             self.event_stream_thread = EventStreamManager(
@@ -285,7 +286,7 @@ def get_environment_flags(self) -> Flags:
         """
         if (
             self.offline_mode or self.enable_local_evaluation
-        ) and self._evaluation_context:
+        ) and self._local_evaluation:
             return self._get_environment_flags_from_document()
         return self._get_environment_flags_from_api()
 
@@ -313,7 +314,7 @@ def get_identity_flags(
         traits = traits or {}
         if (
             self.offline_mode or self.enable_local_evaluation
-        ) and self._evaluation_context:
+        ) and self._local_evaluation:
             return self._get_identity_flags_from_document(identifier, traits)
         return self._get_identity_flags_from_api(
             identifier,
@@ -343,7 +344,7 @@ def get_identity_segments(
         # Identity overrides segments are left out: they aren't segments
         # in Flagsmith, so they're never returned.
         context = map_context_and_identity_data_to_context(
-            context=local_evaluation.base_context,
+            context=local_evaluation.context,
             identifier=identifier,
             traits=traits,
         )
@@ -466,8 +467,8 @@ def update_environment(self) -> None:
             logger.exception("Error retrieving environment document from API")
         else:
             try:
-                self._evaluation_context = map_environment_document_to_context(
-                    environment_data,
+                self._local_evaluation = LocalEvaluation(
+                    map_environment_document_to_context(environment_data)
                 )
                 self._environment_updated_at = (
                     map_environment_document_to_environment_updated_at(
@@ -477,25 +478,6 @@ def update_environment(self) -> None:
             except (KeyError, TypeError, ValueError):
                 logger.exception("Error parsing environment document")
 
-    @property
-    def _evaluation_context(self) -> typing.Optional[SDKEvaluationContext]:
-        return self.__evaluation_context
-
-    @_evaluation_context.setter
-    def _evaluation_context(
-        self, context: typing.Optional[SDKEvaluationContext]
-    ) -> None:
-        """Swap in a new evaluation context and rebuild local evaluation state.
-
-        `LocalEvaluation` holds everything derived from the environment
-        document, e.g. the segment overrides index. Rebuilt once per refresh
-        and reused across every evaluation, without any hot-path cost.
-        """
-        self._local_evaluation = (
-            LocalEvaluation(context) if context is not None else None
-        )
-        self.__evaluation_context = context
-
     def _get_headers(
         self,
         environment_key: str,

flagsmith/local_evaluation.py

@@ -10,7 +10,7 @@
 from flagsmith.mappers import (
     get_identity_overrides_identifiers,
     map_context_and_identity_data_to_context,
-    map_identifiers_to_identity_overrides_rules,
+    map_identifiers_to_identity_overrides_rule,
 )
 from flagsmith.models import (
     Flag,
@@ -68,10 +68,9 @@ class LocalEvaluation:
     """
 
     def __init__(self, context: SDKEvaluationContext) -> None:
-        self.context = context
-
         segments: SegmentContexts = {}
         self._identity_free_segments: SegmentContexts = {}
+        self._identity_overrides_segments: SegmentContexts = {}
         reads_flags = reads_identity = False
         self._identity_overrides_segment_keys: typing.Dict[str, typing.List[str]] = {}
         for segment_key, segment_context in (context.get("segments") or {}).items():
@@ -84,6 +83,7 @@ def __init__(self, context: SDKEvaluationContext) -> None:
                 if not reads.identity:
                     self._identity_free_segments[segment_key] = segment_context
                 continue
+            self._identity_overrides_segments[segment_key] = segment_context
             for identifier in identifiers:
                 self._identity_overrides_segment_keys.setdefault(identifier, []).append(
                     segment_key
@@ -92,16 +92,16 @@ def __init__(self, context: SDKEvaluationContext) -> None:
         self.reads = _ContextReads(flags=reads_flags, identity=reads_identity)
 
         # The context without identity overrides segments.
-        self.base_context: SDKEvaluationContext = {**context, "segments": segments}
+        self.context: SDKEvaluationContext = {**context, "segments": segments}
         self.overrides_index: SegmentOverridesIndex = build_segment_overrides_index(
-            self.base_context
+            self.context
         )
 
     @cached_property
     def environment_flags(self) -> typing.Dict[str, Flag]:
         """Environment flags, which depend on the environment document only."""
         evaluation_result = engine.get_evaluation_result(
-            context={**self.base_context, "segments": self._identity_free_segments},
+            context={**self.context, "segments": self._identity_free_segments},
         )
         return {
             feature_name: Flag.from_evaluation_result(flag_result)
@@ -120,7 +120,7 @@ def get_identity_context(
         many identities share the same overrides.
         """
         context = map_context_and_identity_data_to_context(
-            context=self.base_context,
+            context=self.context,
             identifier=identifier,
             traits=traits,
         )
@@ -129,11 +129,10 @@ def get_identity_context(
 
         identity_segments: SegmentContexts = {}
         identity_overrides_index: SegmentOverridesIndex = {}
-        segment_contexts = self.context.get("segments") or {}
         for segment_key in segment_keys:
             segment_context: SegmentContext[SegmentMetadata, FeatureMetadata] = {
-                **segment_contexts[segment_key],
-                "rules": map_identifiers_to_identity_overrides_rules([identifier]),
+                **self._identity_overrides_segments[segment_key],
+                "rules": [map_identifiers_to_identity_overrides_rule([identifier])],
             }
             identity_segments[segment_key] = segment_context
             for override in segment_context.get("overrides") or ():

flagsmith/mappers.py

@@ -8,6 +8,7 @@
 from flag_engine.context.types import (
     FeatureContext,
     FeatureValue,
+    InOperatorSegmentCondition,
     SegmentContext,
     SegmentRule,
     StrValueSegmentCondition,
@@ -179,7 +180,7 @@ def _map_identity_overrides_to_segments(
         segment_contexts[segment_key] = SegmentContext(
             key=segment_key,
             name="identity_overrides",
-            rules=map_identifiers_to_identity_overrides_rules(identifiers),
+            rules=[map_identifiers_to_identity_overrides_rule(identifiers)],
             overrides=[
                 {
                     "key": "",  # Identity overrides never carry multivariate options
@@ -196,21 +197,15 @@ def _map_identity_overrides_to_segments(
     return segment_contexts
 
 
-def map_identifiers_to_identity_overrides_rules(
+def map_identifiers_to_identity_overrides_rule(
     identifiers: typing.List[str],
-) -> typing.List[SegmentRule]:
-    return [
-        {
-            "type": "ALL",
-            "conditions": [
-                {
-                    "property": "$.identity.identifier",
-                    "operator": "IN",
-                    "value": identifiers,
-                }
-            ],
-        }
-    ]
+) -> SegmentRule:
+    condition: InOperatorSegmentCondition = {
+        "property": "$.identity.identifier",
+        "operator": "IN",
+        "value": identifiers,
+    }
+    return {"type": "ALL", "conditions": [condition]}
 
 
 def get_identity_overrides_identifiers(
@@ -220,10 +215,9 @@ def get_identity_overrides_identifiers(
     any other segment."""
     if (segment_context.get("metadata") or {}).get("source") != "identity_overrides":
         return None
-    # As built by `map_identifiers_to_identity_overrides_rules`.
     [rule] = segment_context["rules"]
     [condition] = rule["conditions"]
-    return typing.cast(typing.List[str], condition["value"])
+    return typing.cast(InOperatorSegmentCondition, condition)["value"]
 
 
 def _map_environment_document_rules_to_context_rules(

flagsmith/models.py

@@ -16,10 +16,12 @@
     SegmentMetadata,
 )
 
-# Segment contexts by their key in the evaluation context.
-SegmentContexts = typing.Dict[str, SegmentContext[SegmentMetadata, FeatureMetadata]]
-# Feature name -> segment contexts overriding that feature.
-SegmentOverridesIndex = typing.MutableMapping[str, SegmentContexts]
+SegmentKey = str
+FeatureName = str
+SegmentContexts = typing.Dict[
+    SegmentKey, SegmentContext[SegmentMetadata, FeatureMetadata]
+]
+SegmentOverridesIndex = typing.MutableMapping[FeatureName, SegmentContexts]
 
 
 def build_segment_overrides_index(

tests/test_flagsmith.py

@@ -17,6 +17,7 @@
     FlagsmithAPIError,
     FlagsmithFeatureDoesNotExistError,
 )
+from flagsmith.local_evaluation import LocalEvaluation
 from flagsmith.mappers import map_environment_document_to_context
 from flagsmith.models import DefaultFlag, ExperimentMetadata, Flag, Flags
 from flagsmith.offline_handlers import OfflineHandler
@@ -48,14 +49,16 @@ def test_update_environment_sets_environment(
 ) -> None:
     # Given
     responses.add(method="GET", url=flagsmith.environment_url, body=environment_json)
-    assert flagsmith._evaluation_context is None
+    assert flagsmith._local_evaluation is None
 
     # When
     flagsmith.update_environment()
 
     # Then
-    assert flagsmith._evaluation_context is not None
-    assert flagsmith._evaluation_context == evaluation_context
+    assert flagsmith._local_evaluation is not None
+    assert flagsmith._local_evaluation.context == (
+        LocalEvaluation(evaluation_context).context
+    )
 
 
 @responses.activate()
@@ -85,7 +88,7 @@ def test_get_environment_flags_uses_local_environment_when_available(
     evaluation_context: SDKEvaluationContext,
 ) -> None:
     # Given
-    flagsmith._evaluation_context = evaluation_context
+    flagsmith._local_evaluation = LocalEvaluation(evaluation_context)
     flagsmith.enable_local_evaluation = True
 
     # When
@@ -150,7 +153,9 @@ def test_get_environment_flags__local_evaluation__includes_identity_free_segment
     # Given
     # the only segment has the given rules, plus identity overrides
     environment["project"]["segments"][0]["rules"] = rules
-    flagsmith._evaluation_context = map_environment_document_to_context(environment)
+    flagsmith._local_evaluation = LocalEvaluation(
+        map_environment_document_to_context(environment)
+    )
     flagsmith.enable_local_evaluation = True
     spy = mocker.spy(engine, "get_evaluation_result")
 
@@ -168,7 +173,6 @@ def test_get_environment_flags__local_evaluation__includes_identity_free_segment
 def test_get_environment_flags__local_evaluation__evaluates_once_per_environment(
     mocker: MockerFixture,
     local_eval_flagsmith: Flagsmith,
-    evaluation_context: SDKEvaluationContext,
 ) -> None:
     # Given
     spy = mocker.spy(engine, "get_evaluation_result")
@@ -183,13 +187,22 @@ def test_get_environment_flags__local_evaluation__evaluates_once_per_environment
     assert first.flags == second.flags
     assert first.flags is not second.flags
 
+
+def test_get_environment_flags__environment_refreshed__evaluates_again(
+    mocker: MockerFixture,
+    local_eval_flagsmith: Flagsmith,
+    evaluation_context: SDKEvaluationContext,
+) -> None:
+    # Given
+    local_eval_flagsmith.get_environment_flags()
+    spy = mocker.spy(engine, "get_evaluation_result")
+
     # When
-    # the environment is refreshed
-    local_eval_flagsmith._evaluation_context = evaluation_context
+    local_eval_flagsmith._local_evaluation = LocalEvaluation(evaluation_context)
     local_eval_flagsmith.get_environment_flags()
 
     # Then
-    assert spy.call_count == 2
+    spy.assert_called_once()
 
 
 @responses.activate()
@@ -254,7 +267,7 @@ def test_get_identity_flags_uses_local_environment_when_available(
     mocker: MockerFixture,
 ) -> None:
     # Given
-    flagsmith._evaluation_context = evaluation_context
+    flagsmith._local_evaluation = LocalEvaluation(evaluation_context)
     flagsmith.enable_local_evaluation = True
     # `Flags` materialises identity flags via `engine.get_evaluation_result`
     # imported from `flagsmith.models`, so patch it where it's actually used.
@@ -350,7 +363,9 @@ def test_get_identity_flags__flags_condition__evaluates_full_context(
             "conditions": [_condition("$.flags.some_feature.enabled")],
         }
     ]
-    flagsmith._evaluation_context = map_environment_document_to_context(environment)
+    flagsmith._local_evaluation = LocalEvaluation(
+        map_environment_document_to_context(environment)
+    )
     flagsmith.enable_local_evaluation = True
     spy = mocker.spy(engine, "get_evaluation_result")