feat: Support generic feature metadata (#269)

Author: khvn26

.gitmodules

@@ -1,4 +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
-	tag = v2.4.0
+	tag = v2.5.0

flag_engine/context/types.py

@@ -11,6 +11,7 @@
 from flag_engine.segments.types import (
     ConditionOperator,
     ContextValue,
+    FeatureMetadataT,
     RuleType,
     SegmentMetadataT,
 )
@@ -54,26 +55,27 @@ class SegmentRule(TypedDict):
     rules: NotRequired[List[SegmentRule]]
 
 
-class FeatureContext(TypedDict):
+class FeatureContext(TypedDict, Generic[FeatureMetadataT]):
     key: str
     feature_key: str
     name: str
     enabled: bool
     value: Any
     variants: NotRequired[List[FeatureValue]]
     priority: NotRequired[float]
+    metadata: NotRequired[FeatureMetadataT]
 
 
-class SegmentContext(TypedDict, Generic[SegmentMetadataT]):
+class SegmentContext(TypedDict, Generic[SegmentMetadataT, FeatureMetadataT]):
     key: str
     name: str
     rules: List[SegmentRule]
-    overrides: NotRequired[List[FeatureContext]]
+    overrides: NotRequired[List[FeatureContext[FeatureMetadataT]]]
     metadata: NotRequired[SegmentMetadataT]
 
 
-class EvaluationContext(TypedDict, Generic[SegmentMetadataT]):
+class EvaluationContext(TypedDict, Generic[SegmentMetadataT, FeatureMetadataT]):
     environment: EnvironmentContext
     identity: NotRequired[Optional[IdentityContext]]
-    segments: NotRequired[Dict[str, SegmentContext[SegmentMetadataT]]]
-    features: NotRequired[Dict[str, FeatureContext]]
+    segments: NotRequired[Dict[str, SegmentContext[SegmentMetadataT, FeatureMetadataT]]]
+    features: NotRequired[Dict[str, FeatureContext[FeatureMetadataT]]]

flag_engine/result/types.py

@@ -8,15 +8,16 @@
 
 from typing_extensions import NotRequired, TypedDict
 
-from flag_engine.segments.types import SegmentMetadataT
+from flag_engine.segments.types import FeatureMetadataT, SegmentMetadataT
 
 
-class FlagResult(TypedDict):
+class FlagResult(TypedDict, Generic[FeatureMetadataT]):
     feature_key: str
     name: str
     enabled: bool
     value: Any
     reason: str
+    metadata: NotRequired[FeatureMetadataT]
 
 
 class SegmentResult(TypedDict, Generic[SegmentMetadataT]):
@@ -25,6 +26,6 @@ class SegmentResult(TypedDict, Generic[SegmentMetadataT]):
     metadata: NotRequired[SegmentMetadataT]
 
 
-class EvaluationResult(TypedDict, Generic[SegmentMetadataT]):
-    flags: Dict[str, FlagResult]
+class EvaluationResult(TypedDict, Generic[SegmentMetadataT, FeatureMetadataT]):
+    flags: Dict[str, FlagResult[FeatureMetadataT]]
     segments: List[SegmentResult[SegmentMetadataT]]

flag_engine/segments/evaluator.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 import operator
 import re
@@ -8,6 +10,7 @@
 
 import jsonpath_rfc9535
 import semver
+from typing_extensions import TypedDict
 
 from flag_engine.context.mappers import map_any_value_to_context_value
 from flag_engine.context.types import (
@@ -23,6 +26,7 @@
 from flag_engine.segments.types import (
     ConditionOperator,
     ContextValue,
+    FeatureMetadataT,
     SegmentMetadataT,
     is_context_value,
 )
@@ -32,24 +36,27 @@
 from flag_engine.utils.types import SupportsStr, get_casting_function
 
 
-class FeatureContextWithSegmentName(typing.TypedDict):
-    feature_context: FeatureContext
+class FeatureContextWithSegmentName(TypedDict, typing.Generic[FeatureMetadataT]):
+    feature_context: FeatureContext[FeatureMetadataT]
     segment_name: str
 
 
 def get_evaluation_result(
-    context: EvaluationContext[SegmentMetadataT],
-) -> EvaluationResult[SegmentMetadataT]:
+    context: EvaluationContext[SegmentMetadataT, FeatureMetadataT],
+) -> EvaluationResult[SegmentMetadataT, FeatureMetadataT]:
     """
     Get the evaluation result for a given context.
 
     :param context: the evaluation context
     :return: EvaluationResult containing the context, flags, and segments
     """
     segments: list[SegmentResult[SegmentMetadataT]] = []
-    flags: dict[str, FlagResult] = {}
+    flags: dict[str, FlagResult[FeatureMetadataT]] = {}
 
-    segment_feature_contexts: dict[SupportsStr, FeatureContextWithSegmentName] = {}
+    segment_feature_contexts: dict[
+        SupportsStr,
+        FeatureContextWithSegmentName[FeatureMetadataT],
+    ] = {}
 
     for segment_context in (context.get("segments") or {}).values():
         if not is_context_in_segment(context, segment_context):
@@ -59,8 +66,8 @@ def get_evaluation_result(
             "key": segment_context["key"],
             "name": segment_context["name"],
         }
-        if metadata := segment_context.get("metadata"):
-            segment_result["metadata"] = metadata
+        if segment_metadata := segment_context.get("metadata"):
+            segment_result["metadata"] = segment_metadata
         segments.append(segment_result)
 
         if overrides := segment_context.get("overrides"):
@@ -95,13 +102,16 @@ def get_evaluation_result(
             feature_context["feature_key"],
         ):
             feature_context = feature_context_with_segment_name["feature_context"]
-            flags[feature_name] = {
+            flag_result: FlagResult[FeatureMetadataT]
+            flags[feature_name] = flag_result = {
                 "enabled": feature_context["enabled"],
                 "feature_key": feature_context["feature_key"],
                 "name": feature_context["name"],
                 "reason": f"TARGETING_MATCH; segment={feature_context_with_segment_name['segment_name']}",
                 "value": feature_context.get("value"),
             }
+            if feature_metadata := feature_context.get("metadata"):
+                flag_result["metadata"] = feature_metadata
             continue
         flags[feature_name] = get_flag_result_from_feature_context(
             feature_context=feature_context,
@@ -115,9 +125,9 @@ def get_evaluation_result(
 
 
 def get_flag_result_from_feature_context(
-    feature_context: FeatureContext,
+    feature_context: FeatureContext[FeatureMetadataT],
     key: typing.Optional[SupportsStr],
-) -> FlagResult:
+) -> FlagResult[FeatureMetadataT]:
     """
     Get a feature value from the feature context
     for a given key.
@@ -126,6 +136,8 @@ def get_flag_result_from_feature_context(
     :param key: the key to get the value for
     :return: the value for the key in the feature context
     """
+    flag_result: typing.Optional[FlagResult[FeatureMetadataT]] = None
+
     if key is not None and (variants := feature_context.get("variants")):
         percentage_value = get_hashed_percentage_for_object_ids(
             [feature_context["key"], key]
@@ -139,28 +151,35 @@ def get_flag_result_from_feature_context(
         ):
             limit = (weight := variant["weight"]) + start_percentage
             if start_percentage <= percentage_value < limit:
-                return {
+                flag_result = {
                     "enabled": feature_context["enabled"],
                     "feature_key": feature_context["feature_key"],
                     "name": feature_context["name"],
                     "reason": f"SPLIT; weight={weight}",
                     "value": variant["value"],
                 }
+                break
 
             start_percentage = limit
 
-    return {
-        "enabled": feature_context["enabled"],
-        "feature_key": feature_context["feature_key"],
-        "name": feature_context["name"],
-        "reason": "DEFAULT",
-        "value": feature_context["value"],
-    }
+    if flag_result is None:
+        flag_result = {
+            "enabled": feature_context["enabled"],
+            "feature_key": feature_context["feature_key"],
+            "name": feature_context["name"],
+            "reason": "DEFAULT",
+            "value": feature_context["value"],
+        }
+
+    if metadata := feature_context.get("metadata"):
+        flag_result["metadata"] = metadata
+
+    return flag_result
 
 
 def is_context_in_segment(
-    context: EvaluationContext[SegmentMetadataT],
-    segment_context: SegmentContext[SegmentMetadataT],
+    context: EvaluationContext[typing.Any, typing.Any],
+    segment_context: SegmentContext[typing.Any, typing.Any],
 ) -> bool:
     return bool(rules := segment_context["rules"]) and all(
         context_matches_rule(
@@ -171,7 +190,7 @@ def is_context_in_segment(
 
 
 def context_matches_rule(
-    context: EvaluationContext[SegmentMetadataT],
+    context: EvaluationContext[typing.Any, typing.Any],
     rule: SegmentRule,
     segment_key: SupportsStr,
 ) -> bool:
@@ -201,7 +220,7 @@ def context_matches_rule(
 
 
 def context_matches_condition(
-    context: EvaluationContext[SegmentMetadataT],
+    context: EvaluationContext[typing.Any, typing.Any],
     condition: SegmentCondition,
     segment_key: SupportsStr,
 ) -> bool:
@@ -262,7 +281,7 @@ def context_matches_condition(
 
 
 def get_context_value(
-    context: EvaluationContext[SegmentMetadataT],
+    context: EvaluationContext[typing.Any, typing.Any],
     property: str,
 ) -> ContextValue:
     value = None

flag_engine/segments/types.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Literal, Union, get_args
+from typing import Any, Literal, Mapping, Union, get_args
 
 from typing_extensions import TypeGuard, TypeVar
 
-SegmentMetadataT = TypeVar("SegmentMetadataT", default=Dict[str, object])
+SegmentMetadataT = TypeVar("SegmentMetadataT", default=Mapping[str, object])
+FeatureMetadataT = TypeVar("FeatureMetadataT", default=Mapping[str, object])
 
 ConditionOperator = Literal[
     "EQUAL",