feat: Context values for Segments

Author: khvn26

flag_engine/context/__init__.py

@@ -0,0 +1,29 @@
+from flag_engine.context.types import EvaluationContext
+from flag_engine.environments.models import EnvironmentModel
+from flag_engine.identities.models import IdentityModel
+
+
+def map_environment_identity_to_context(
+    environment: EnvironmentModel,
+    identity: IdentityModel,
+) -> EvaluationContext:
+    """
+    Maps an EnvironmentModel and IdentityModel to an EvaluationContext.
+
+    :param environment: The environment model object.
+    :param identity: The identity model object.
+    :return: An EvaluationContext containing the environment and identity.
+    """
+    return {
+        "environment": {
+            "key": environment.api_key,
+            "name": environment.name or "",
+        },
+        "identity": {
+            "identifier": identity.identifier,
+            "key": str(identity.django_id or identity.composite_key),
+            "traits": {
+                trait.trait_key: trait.trait_value for trait in identity.identity_traits
+            },
+        },
+    }

flag_engine/context/mappers.py

@@ -0,0 +1,28 @@
+# generated by datamodel-codegen:
+#   filename:  https://raw.githubusercontent.com/Flagsmith/flagsmith/chore/update-evaluation-context/sdk/evaluation-context.json  # noqa: E501
+#   timestamp: 2025-07-16T10:39:10+00:00
+
+from __future__ import annotations
+
+from typing import Dict, Optional, TypedDict
+
+from typing_extensions import NotRequired
+
+from flag_engine.identities.traits.types import TraitValue
+from flag_engine.utils.types import SupportsStr
+
+
+class EnvironmentContext(TypedDict):
+    key: str
+    name: str
+
+
+class IdentityContext(TypedDict):
+    identifier: str
+    key: SupportsStr
+    traits: NotRequired[Dict[str, TraitValue]]
+
+
+class EvaluationContext(TypedDict):
+    environment: EnvironmentContext
+    identity: NotRequired[Optional[IdentityContext]]

flag_engine/context/types.py

@@ -1,10 +1,11 @@
 import typing
 
+from flag_engine.context.mappers import map_environment_identity_to_context
 from flag_engine.environments.models import EnvironmentModel
 from flag_engine.features.models import FeatureModel, FeatureStateModel
 from flag_engine.identities.models import IdentityModel
 from flag_engine.identities.traits.models import TraitModel
-from flag_engine.segments.evaluator import get_identity_segments
+from flag_engine.segments.evaluator import get_context_segments
 from flag_engine.utils.exceptions import FeatureStateNotFound
 
 
@@ -99,26 +100,43 @@ def _get_identity_feature_states_dict(
     override_traits: typing.Optional[typing.List[TraitModel]],
 ) -> typing.Dict[FeatureModel, FeatureStateModel]:
     # Get feature states from the environment
-    feature_states = {fs.feature: fs for fs in environment.feature_states}
+    feature_states_by_feature = {fs.feature: fs for fs in environment.feature_states}
+
+    context = map_environment_identity_to_context(
+        environment=environment,
+        identity=identity,
+    )
+    if override_traits:
+        if typing.TYPE_CHECKING:
+            assert context["identity"]
+        context["identity"].setdefault("traits", {}).update(
+            {trait.trait_key: trait.trait_value for trait in override_traits}
+        )
 
     # Override with any feature states defined by matching segments
-    identity_segments = get_identity_segments(environment, identity, override_traits)
-    for matching_segment in identity_segments:
-        for feature_state in matching_segment.feature_states:
-            if feature_state.feature in feature_states:
-                if feature_states[feature_state.feature].is_higher_segment_priority(
-                    feature_state
-                ):
-                    continue
-            feature_states[feature_state.feature] = feature_state
+    for context_segment in get_context_segments(
+        context=context,
+        segments=environment.project.segments,
+    ):
+        for segment_feature_state in context_segment.feature_states:
+            if (
+                environment_feature_state := feature_states_by_feature.get(
+                    segment_feature := segment_feature_state.feature
+                )
+            ) and environment_feature_state.is_higher_segment_priority(
+                segment_feature_state
+            ):
+                continue
+            feature_states_by_feature[segment_feature] = segment_feature_state
 
     # Override with any feature states defined directly the identity
-    feature_states.update(
+    feature_states_by_feature.update(
         {
-            fs.feature: fs
-            for fs in identity.identity_features
-            if fs.feature in feature_states
+            identity_feature: identity_feature_state
+            for identity_feature_state in identity.identity_features
+            if (identity_feature := identity_feature_state.feature)
+            in feature_states_by_feature
         }
     )
 
-    return feature_states
+    return feature_states_by_feature