chore: replacing-deprecated-methods

flagsmith/flagsmith.py

@@ -10,13 +10,13 @@
 from flag_engine.identities.models import IdentityModel
 from flag_engine.identities.traits.models import TraitModel
 from flag_engine.identities.traits.types import TraitValue
-from flag_engine.segments.evaluator import get_identity_segments
 from requests.adapters import HTTPAdapter
 from requests.utils import default_user_agent
 from urllib3 import Retry
 
 from flagsmith.analytics import AnalyticsProcessor
 from flagsmith.exceptions import FlagsmithAPIError, FlagsmithClientError
+from flagsmith.mappers import map_environment_identity_to_context
 from flagsmith.models import DefaultFlag, Flags, Segment
 from flagsmith.offline_handlers import BaseOfflineHandler
 from flagsmith.polling_manager import EnvironmentDataPollingManager
@@ -280,10 +280,18 @@ def get_identity_segments(
 
         traits = traits or {}
         identity_model = self._get_identity_model(identifier, **traits)
-        segment_models = get_identity_segments(
-            environment=self._environment, identity=identity_model
+        context = map_environment_identity_to_context(
+            environment=self._environment,
+            identity=identity_model,
+            override_traits=None,
         )
-        return [Segment(id=sm.id, name=sm.name) for sm in segment_models]
+        evaluation_result = engine.get_evaluation_result(
+            context=context,
+        )
+        return [
+            Segment(id=int(sm["key"]), name=sm["name"])
+            for sm in evaluation_result.get("segments", [])
+        ]
 
     def update_environment(self) -> None:
         try:
@@ -321,8 +329,18 @@ def _get_environment_from_api(self) -> EnvironmentModel:
     def _get_environment_flags_from_document(self) -> Flags:
         if self._environment is None:
             raise TypeError("No environment present")
-        return Flags.from_feature_state_models(
-            feature_states=engine.get_environment_feature_states(self._environment),
+        identity = self._get_identity_model(identifier="", traits=None)
+
+        context = map_environment_identity_to_context(
+            environment=self._environment,
+            identity=identity,
+            override_traits=None,
+        )
+
+        evaluation_result = engine.get_evaluation_result(context=context)
+
+        return Flags.from_evaluation_result(
+            evaluation_result=evaluation_result,
             analytics_processor=self._analytics_processor,
             default_flag_handler=self.default_flag_handler,
         )
@@ -333,13 +351,20 @@ def _get_identity_flags_from_document(
         identity_model = self._get_identity_model(identifier, **traits)
         if self._environment is None:
             raise TypeError("No environment present")
-        feature_states = engine.get_identity_feature_states(
-            self._environment, identity_model
+
+        context = map_environment_identity_to_context(
+            environment=self._environment,
+            identity=identity_model,
+            override_traits=None,
         )
-        return Flags.from_feature_state_models(
-            feature_states=feature_states,
+
+        evaluation_result = engine.get_evaluation_result(
+            context=context,
+        )
+
+        return Flags.from_evaluation_result(
+            evaluation_result=evaluation_result,
             analytics_processor=self._analytics_processor,
-            identity_id=identity_model.composite_key,
             default_flag_handler=self.default_flag_handler,
         )
 

flagsmith/mappers.py

@@ -0,0 +1,209 @@
+import typing
+from collections import defaultdict
+
+from flag_engine.context.types import (
+    EvaluationContext,
+    FeatureContext,
+    SegmentContext,
+    SegmentRule,
+)
+from flag_engine.environments.models import EnvironmentModel
+from flag_engine.features.models import (
+    FeatureStateModel,
+    MultivariateFeatureStateValueModel,
+)
+from flag_engine.identities.models import IdentityModel
+from flag_engine.identities.traits.models import TraitModel
+from flag_engine.segments.models import SegmentRuleModel
+
+OverrideKey = typing.Tuple[
+    str,
+    str,
+    bool,
+    typing.Any,
+]
+OverridesKey = typing.Tuple[OverrideKey, ...]
+
+
+def map_environment_identity_to_context(
+    environment: EnvironmentModel,
+    identity: IdentityModel,
+    override_traits: typing.Optional[typing.List[TraitModel]],
+) -> EvaluationContext:
+    """
+    Map 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.
+    """
+    features = map_feature_states_to_feature_contexts(environment.feature_states)
+    segments: typing.Dict[str, SegmentContext] = {}
+    for segment in environment.project.segments:
+        segment_ctx_data: SegmentContext = {
+            "key": str(segment.id),
+            "name": segment.name,
+            "rules": map_segment_rules_to_segment_context_rules(segment.rules),
+        }
+        if segment_feature_states := segment.feature_states:
+            segment_ctx_data["overrides"] = list(
+                map_feature_states_to_feature_contexts(segment_feature_states).values()
+            )
+        segments[segment.name] = segment_ctx_data
+    # Concatenate feature states overriden for identities
+    # to segment contexts
+    features_to_identifiers: typing.Dict[
+        OverridesKey,
+        typing.List[str],
+    ] = defaultdict(list)
+    for identity_override in (*environment.identity_overrides, identity):
+        identity_features: typing.List[FeatureStateModel] = (
+            identity_override.identity_features
+        )
+        if not identity_features:
+            continue
+        overrides_key = tuple(
+            (
+                str(feature_state.feature.id),
+                feature_state.feature.name,
+                feature_state.enabled,
+                feature_state.feature_state_value,
+            )
+            for feature_state in sorted(identity_features, key=_get_name)
+        )
+        features_to_identifiers[overrides_key].append(identity_override.identifier)
+    for overrides_key, identifiers in features_to_identifiers.items():
+        segment_name = f"overrides_{abs(hash(overrides_key))}"
+        segments[segment_name] = SegmentContext(
+            key="",  # Identity override segments never use % Split operator
+            name=segment_name,
+            rules=[
+                {
+                    "type": "ALL",
+                    "rules": [
+                        {
+                            "type": "ALL",
+                            "conditions": [
+                                {
+                                    "property": "$.identity.identifier",
+                                    "operator": "IN",
+                                    "value": ",".join(identifiers),
+                                }
+                            ],
+                        }
+                    ],
+                }
+            ],
+            overrides=[
+                {
+                    "key": "",  # Identity overrides never carry multivariate options
+                    "feature_key": feature_key,
+                    "name": feature_name,
+                    "enabled": feature_enabled,
+                    "value": feature_value,
+                    "priority": float("-inf"),  # Highest possible priority
+                }
+                for feature_key, feature_name, feature_enabled, feature_value in overrides_key
+            ],
+        )
+    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
+                )
+            },
+        },
+        "features": features,
+        "segments": segments,
+    }
+
+
+def map_feature_states_to_feature_contexts(
+    feature_states: typing.List[FeatureStateModel],
+) -> typing.Dict[str, FeatureContext]:
+    """
+    Map feature states to feature contexts.
+
+    :param feature_states: A list of FeatureStateModel objects.
+    :return: A dictionary mapping feature names to their contexts.
+    """
+    features: typing.Dict[str, FeatureContext] = {}
+    for feature_state in feature_states:
+        feature_ctx_data: FeatureContext = {
+            "key": str(feature_state.django_id or feature_state.featurestate_uuid),
+            "feature_key": str(feature_state.feature.id),
+            "name": feature_state.feature.name,
+            "enabled": feature_state.enabled,
+            "value": feature_state.feature_state_value,
+        }
+        multivariate_feature_state_values: typing.List[
+            MultivariateFeatureStateValueModel
+        ]
+        if (
+            multivariate_feature_state_values := feature_state.multivariate_feature_state_values
+        ):
+            feature_ctx_data["variants"] = [
+                {
+                    "value": multivariate_feature_state_value.multivariate_feature_option.value,
+                    "weight": multivariate_feature_state_value.percentage_allocation,
+                }
+                for multivariate_feature_state_value in sorted(
+                    multivariate_feature_state_values,
+                    key=_get_multivariate_feature_state_value_id,
+                )
+            ]
+        if feature_segment := feature_state.feature_segment:
+            if (priority := feature_segment.priority) is not None:
+                feature_ctx_data["priority"] = priority
+        features[feature_state.feature.name] = feature_ctx_data
+    return features
+
+
+def _get_multivariate_feature_state_value_id(
+    multivariate_feature_state_value: MultivariateFeatureStateValueModel,
+) -> int:
+    return (
+        multivariate_feature_state_value.id
+        or multivariate_feature_state_value.mv_fs_value_uuid.int
+    )
+
+
+def map_segment_rules_to_segment_context_rules(
+    rules: typing.List[SegmentRuleModel],
+) -> typing.List[SegmentRule]:
+    """
+    Map segment rules to segment rules for the evaluation context.
+
+    :param rules: A list of SegmentRuleModel objects.
+    :return: A list of SegmentRule objects.
+    """
+    return [
+        {
+            "type": rule.type,
+            "conditions": [
+                {
+                    "property": condition.property_ or "",
+                    "operator": condition.operator,
+                    "value": condition.value or "",
+                }
+                for condition in rule.conditions
+            ],
+            "rules": map_segment_rules_to_segment_context_rules(rule.rules),
+        }
+        for rule in rules
+    ]
+
+
+def _get_name(feature_state: FeatureStateModel) -> str:
+    return feature_state.feature.name

flagsmith/models.py

@@ -3,7 +3,7 @@
 import typing
 from dataclasses import dataclass, field
 
-from flag_engine.features.models import FeatureStateModel
+from flag_engine.result.types import EvaluationResult, FlagResult
 
 from flagsmith.analytics import AnalyticsProcessor
 from flagsmith.exceptions import FlagsmithFeatureDoesNotExistError
@@ -27,16 +27,15 @@ class Flag(BaseFlag):
     is_default: bool = field(default=False)
 
     @classmethod
-    def from_feature_state_model(
+    def from_evaluation_result(
         cls,
-        feature_state_model: FeatureStateModel,
-        identity_id: typing.Optional[typing.Union[str, int]] = None,
+        flag: FlagResult,
     ) -> Flag:
         return Flag(
-            enabled=feature_state_model.enabled,
-            value=feature_state_model.get_value(identity_id=identity_id),
-            feature_name=feature_state_model.feature.name,
-            feature_id=feature_state_model.feature.id,
+            enabled=flag["enabled"],
+            value=flag["value"],
+            feature_name=flag["name"],
+            feature_id=int(flag["feature_key"]),
         )
 
     @classmethod
@@ -56,22 +55,22 @@ class Flags:
     _analytics_processor: typing.Optional[AnalyticsProcessor] = None
 
     @classmethod
-    def from_feature_state_models(
+    def from_evaluation_result(
         cls,
-        feature_states: typing.Sequence[FeatureStateModel],
+        evaluation_result: EvaluationResult,
         analytics_processor: typing.Optional[AnalyticsProcessor],
         default_flag_handler: typing.Optional[typing.Callable[[str], DefaultFlag]],
-        identity_id: typing.Optional[typing.Union[str, int]] = None,
     ) -> Flags:
-        flags = {
-            feature_state.feature.name: Flag.from_feature_state_model(
-                feature_state, identity_id=identity_id
-            )
-            for feature_state in feature_states
-        }
-
         return cls(
-            flags=flags,
+            flags={
+                flag["name"]: Flag(
+                    enabled=flag["enabled"],
+                    value=flag["value"],
+                    feature_name=flag["name"],
+                    feature_id=int(flag["feature_key"]),
+                )
+                for flag in evaluation_result["flags"]
+            },
             default_flag_handler=default_flag_handler,
             _analytics_processor=analytics_processor,
         )