feat: Context values for Segments (#220)

Author: khvn26

.gitmodules

@@ -1,3 +1,4 @@
 [submodule "tests/engine_tests/engine-test-data"]
 	path = tests/engine_tests/engine-test-data
 	url = https://github.com/flagsmith/engine-test-data.git
+	branch = feat/context-values

flag_engine/context/__init__.py

@@ -0,0 +1,39 @@
+import typing
+
+from flag_engine.context.types import EvaluationContext
+from flag_engine.environments.models import EnvironmentModel
+from flag_engine.identities.models import IdentityModel
+from flag_engine.identities.traits.models import TraitModel
+
+
+def map_environment_identity_to_context(
+    environment: EnvironmentModel,
+    identity: IdentityModel,
+    override_traits: typing.Optional[typing.List[TraitModel]],
+) -> EvaluationContext:
+    """
+    Maps an EnvironmentModel and IdentityModel to an EvaluationContext.
+
+    :param environment: The environment model object.
+    :param identity: The identity model object.
+    :param override_traits: A list of TraitModel objects, to be used in place of `identity.identity_traits` if provided.
+    :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 (
+                    override_traits
+                    if override_traits is not None
+                    else 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 ContextValue
+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, ContextValue]]
+
+
+class EvaluationContext(TypedDict):
+    environment: EnvironmentContext
+    identity: NotRequired[Optional[IdentityContext]]

flag_engine/context/types.py

@@ -1,10 +1,12 @@
 import typing
 
+from flag_engine.context.mappers import map_environment_identity_to_context
+from flag_engine.context.types import EvaluationContext
 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
 
 
@@ -53,9 +55,17 @@ def get_identity_feature_states(
     :return: list of feature state models based on the environment, any matching
         segments and any specific identity overrides
     """
+    context = map_environment_identity_to_context(
+        environment=environment,
+        identity=identity,
+        override_traits=override_traits,
+    )
+
     feature_states = list(
         _get_identity_feature_states_dict(
-            environment, identity, override_traits
+            environment=environment,
+            identity=identity,
+            context=context,
         ).values()
     )
     if environment.get_hide_disabled_flags():
@@ -79,8 +89,16 @@ def get_identity_feature_state(
     :return: feature state model based on the environment, any matching
         segments and any specific identity overrides
     """
+    context = map_environment_identity_to_context(
+        environment=environment,
+        identity=identity,
+        override_traits=override_traits,
+    )
+
     feature_states = _get_identity_feature_states_dict(
-        environment, identity, override_traits
+        environment=environment,
+        identity=identity,
+        context=context,
     )
     matching_feature = next(
         filter(lambda feature: feature.name == feature_name, feature_states.keys()),
@@ -96,29 +114,33 @@ def get_identity_feature_state(
 def _get_identity_feature_states_dict(
     environment: EnvironmentModel,
     identity: IdentityModel,
-    override_traits: typing.Optional[typing.List[TraitModel]],
+    context: EvaluationContext,
 ) -> 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}
 
     # 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 (
+                feature_state := feature_states_by_feature.get(
+                    segment_feature := segment_feature_state.feature
+                )
+            ) and 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

flag_engine/engine.py

@@ -1,8 +1,8 @@
 from pydantic import BaseModel, Field
 
-from flag_engine.identities.traits.types import TraitValue
+from flag_engine.identities.traits.types import ContextValue
 
 
 class TraitModel(BaseModel):
     trait_key: str
-    trait_value: TraitValue = Field(...)
+    trait_value: ContextValue = Field(...)

flag_engine/identities/traits/models.py

@@ -8,10 +8,10 @@
 
 from flag_engine.identities.traits.constants import TRAIT_STRING_VALUE_MAX_LENGTH
 
-_UnconstrainedTraitValue = Union[None, int, float, bool, str]
+_UnconstrainedContextValue = Union[None, int, float, bool, str]
 
 
-def map_any_value_to_trait_value(value: Any) -> _UnconstrainedTraitValue:
+def map_any_value_to_trait_value(value: Any) -> _UnconstrainedContextValue:
     """
     Try to coerce a value of arbitrary type to a trait value type.
     Union member-specific constraints, such as max string value length, are ignored here.
@@ -36,19 +36,19 @@ def map_any_value_to_trait_value(value: Any) -> _UnconstrainedTraitValue:
 _float_pattern = re.compile(r"-?[0-9]+\.[0-9]+")
 
 
-def _map_string_value_to_trait_value(value: str) -> _UnconstrainedTraitValue:
+def _map_string_value_to_trait_value(value: str) -> _UnconstrainedContextValue:
     if _int_pattern.fullmatch(value):
         return int(value)
     if _float_pattern.fullmatch(value):
         return float(value)
     return value
 
 
-def _is_trait_value(value: Any) -> TypeGuard[_UnconstrainedTraitValue]:
-    return isinstance(value, get_args(_UnconstrainedTraitValue))
+def _is_trait_value(value: Any) -> TypeGuard[_UnconstrainedContextValue]:
+    return isinstance(value, get_args(_UnconstrainedContextValue))
 
 
-TraitValue = Annotated[
+ContextValue = Annotated[
     Union[
         None,
         StrictBool,