perf: Map identity overrides, evaluate environment flags once per document refresh

flagsmith/flagsmith.py

@@ -15,28 +15,20 @@
     EventProcessorConfig,
 )
 from flagsmith.exceptions import FlagsmithAPIError, FlagsmithClientError
+from flagsmith.local_evaluation import LocalEvaluation
 from flagsmith.mappers import (
     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,
     resolve_trait_values,
 )
-from flagsmith.models import (
-    DefaultFlag,
-    Flag,
-    Flags,
-    Segment,
-    SegmentOverridesIndex,
-    build_segment_overrides_index,
-)
+from flagsmith.models import DefaultFlag, Flag, Flags, Segment
 from flagsmith.offline_handlers import OfflineHandler
 from flagsmith.polling_manager import EnvironmentDataPollingManager
 from flagsmith.streaming_manager import EventStreamManager
 from flagsmith.types import (
     ApplicationMetadata,
     JsonType,
-    SDKEvaluationContext,
     StreamEvent,
     TraitMapping,
 )
@@ -135,8 +127,7 @@ 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._segment_overrides_index: SegmentOverridesIndex = {}
+        self._local_evaluation: typing.Optional[LocalEvaluation] = None
         self._environment_updated_at: typing.Optional[datetime] = None
 
         # argument validation
@@ -158,7 +149,7 @@ def __init__(
             )
 
         if self.offline_handler:
-            self._evaluation_context = map_environment_document_to_context(
+            self._local_evaluation = LocalEvaluation.from_environment_document(
                 self.offline_handler.get_environment()
             )
 
@@ -248,12 +239,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(
@@ -291,7 +283,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()
 
@@ -319,7 +311,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,
@@ -341,13 +333,13 @@ def get_identity_segments(
             Flagsmith, e.g. {"num_orders": 10}
         :return: list of Segment objects that the identity is part of.
         """
-        if not self._evaluation_context:
+        if not (local_evaluation := self._local_evaluation):
             raise FlagsmithClientError(
                 "Local evaluation required to obtain identity segments."
             )
 
         context = map_context_and_identity_data_to_context(
-            context=self._evaluation_context,
+            context=local_evaluation.context,
             identifier=identifier,
             traits=traits,
         )
@@ -470,8 +462,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.from_environment_document(
+                    environment_data
                 )
                 self._environment_updated_at = (
                     map_environment_document_to_environment_updated_at(
@@ -481,26 +473,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 the overrides index.
-
-        The index maps feature_name -> segments that override it. Built once
-        per refresh and reused across every subsequent per-identity lazy
-        resolution; rebuilding here keeps it in sync with the current doc
-        without any hot-path cost.
-        """
-        self.__evaluation_context = context
-        self._segment_overrides_index = (
-            build_segment_overrides_index(context) if context is not None else {}
-        )
-
     def _get_headers(
         self,
         environment_key: str,
@@ -520,34 +492,24 @@ def _get_headers(
         return headers
 
     def _get_environment_flags_from_document(self) -> Flags:
-        if self._evaluation_context is None:
+        if (local_evaluation := self._local_evaluation) is None:
             raise TypeError("No environment present")
 
-        # Omit segments from evaluation context for environment flags
-        # as they are only relevant for identity-specific evaluations
-        context_without_segments = self._evaluation_context.copy()
-        context_without_segments.pop("segments", None)
-
-        evaluation_result = engine.get_evaluation_result(
-            context=context_without_segments,
-        )
-
-        return Flags.from_evaluation_result(
-            evaluation_result=evaluation_result,
-            analytics_processor=self._analytics_processor,
+        return Flags(
+            flags=dict(local_evaluation.environment_flags),
             default_flag_handler=self.default_flag_handler,
+            _analytics_processor=self._analytics_processor,
         )
 
     def _get_identity_flags_from_document(
         self,
         identifier: str,
         traits: TraitMapping,
     ) -> Flags:
-        if self._evaluation_context is None:
+        if (local_evaluation := self._local_evaluation) is None:
             raise TypeError("No environment present")
 
-        context = map_context_and_identity_data_to_context(
-            context=self._evaluation_context,
+        context, overrides_index = local_evaluation.get_identity_context(
             identifier=identifier,
             traits=traits,
         )
@@ -556,7 +518,7 @@ def _get_identity_flags_from_document(
         # large environment.
         return Flags.from_evaluation_context(
             context=context,
-            overrides_index=self._segment_overrides_index,
+            overrides_index=overrides_index,
             analytics_processor=self._analytics_processor,
             default_flag_handler=self.default_flag_handler,
         )

flagsmith/local_evaluation.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import typing
+from collections import ChainMap
+from functools import cached_property
+
+from flag_engine import engine
+from flag_engine.context.types import SegmentContext
+
+from flagsmith.api.types import EnvironmentModel
+from flagsmith.mappers import (
+    map_context_and_identity_data_to_context,
+    map_environment_document_to_context,
+    map_environment_document_to_identity_overrides_segments,
+)
+from flagsmith.models import (
+    Flag,
+    SegmentOverridesIndex,
+    build_segment_overrides_index,
+)
+from flagsmith.types import (
+    FeatureMetadata,
+    SDKEvaluationContext,
+    SegmentMetadata,
+    TraitMapping,
+)
+
+
+class LocalEvaluation:
+    """Per-environment-document state for local evaluation."""
+
+    def __init__(
+        self,
+        context: SDKEvaluationContext,
+        identity_overrides_segments: typing.Optional[
+            typing.Mapping[str, SegmentContext[SegmentMetadata, FeatureMetadata]]
+        ] = None,
+    ) -> None:
+        self.context = context
+        self.overrides_index: SegmentOverridesIndex = build_segment_overrides_index(
+            context
+        )
+        self._identity_overrides = {
+            identifier: (segment_context, self._index_overrides(segment_context))
+            for identifier, segment_context in (
+                identity_overrides_segments or {}
+            ).items()
+        }
+
+    @classmethod
+    def from_environment_document(
+        cls,
+        environment_document: EnvironmentModel,
+    ) -> LocalEvaluation:
+        return cls(
+            context=map_environment_document_to_context(environment_document),
+            identity_overrides_segments=(
+                map_environment_document_to_identity_overrides_segments(
+                    environment_document
+                )
+            ),
+        )
+
+    @cached_property
+    def environment_flags(self) -> typing.Dict[str, Flag]:
+        """Environment flags, which depend on the environment document only."""
+        context_without_segments = self.context.copy()
+        context_without_segments.pop("segments", None)
+        evaluation_result = engine.get_evaluation_result(
+            context=context_without_segments,
+        )
+        return {
+            feature_name: Flag.from_evaluation_result(flag_result)
+            for feature_name, flag_result in evaluation_result["flags"].items()
+        }
+
+    def get_identity_context(
+        self,
+        identifier: str,
+        traits: typing.Optional[TraitMapping] = None,
+    ) -> typing.Tuple[SDKEvaluationContext, SegmentOverridesIndex]:
+        """Evaluation context and overrides index for an identity, including
+        its own identity overrides segment, if any."""
+        context = map_context_and_identity_data_to_context(
+            context=self.context,
+            identifier=identifier,
+            traits=traits,
+        )
+        if (identity_overrides := self._identity_overrides.get(identifier)) is None:
+            return context, self.overrides_index
+
+        segment_context, identity_overrides_index = identity_overrides
+        context["segments"] = {
+            **(context.get("segments") or {}),
+            segment_context["key"]: segment_context,
+        }
+        return context, ChainMap(identity_overrides_index, self.overrides_index)
+
+    def _index_overrides(
+        self,
+        segment_context: SegmentContext[SegmentMetadata, FeatureMetadata],
+    ) -> SegmentOverridesIndex:
+        """The overrides index, for the features `segment_context` overrides."""
+        return {
+            feature_name: {
+                **(self.overrides_index.get(feature_name) or {}),
+                segment_context["key"]: segment_context,
+            }
+            for override in segment_context.get("overrides") or ()
+            if (feature_name := override["name"])
+        }