optimise for traits

Author: khvn26

flag_engine/segments/evaluator.py

@@ -261,7 +261,12 @@ def get_context_value(
     context: EvaluationContext,
     property: str,
 ) -> ContextValue:
-    return _get_context_value_getter(property)(context)
+    if property.startswith("$."):
+        return _get_context_value_getter(property)(context)
+    if identity_context := context.get("identity"):
+        if traits := identity_context.get("traits"):
+            return traits.get(property)
+    return None
 
 
 def _matches_context_value(
@@ -385,6 +390,8 @@ def _get_context_value_getter(
     try:
         compiled_query = jsonpath_rfc9535.compile(property)
     except jsonpath_rfc9535.JSONPathSyntaxError:
+        # This covers a rare case when a trait starting with "$.",
+        # but not a valid JSONPath, is used.
         compiled_query = jsonpath_rfc9535.compile(
             f'$.identity.traits["{escape_double_quotes(property)}"]',
         )

tests/unit/segments/test_segments_evaluator.py

@@ -1048,3 +1048,19 @@ def test_get_context_value__invalid_jsonpath__returns_expected(
 
     # Then
     assert result is None
+
+
+def test_get_context_value__jsonpath_like_trait__returns_expected(
+    context: EvaluationContext,
+) -> None:
+    # Given
+    jsonpath_like_trait = '$.  i am not" a valid jsonpath'
+    expected_result = "some_value"
+    assert context["identity"]
+    context["identity"]["traits"][jsonpath_like_trait] = expected_result
+
+    # When
+    result = get_context_value(context, jsonpath_like_trait)
+
+    # Then
+    assert result == expected_result