fix: $.-prefixed traits are not prioritised over valid JSONPaths (#266)

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.1.0
+	tag = v2.2.0

flag_engine/segments/evaluator.py

@@ -4,7 +4,7 @@
 import typing
 import warnings
 from contextlib import suppress
-from functools import lru_cache, wraps
+from functools import lru_cache, partial, wraps
 
 import jsonpath_rfc9535
 import semver
@@ -26,7 +26,7 @@
     SegmentMetadataT,
     is_context_value,
 )
-from flag_engine.segments.utils import escape_double_quotes, get_matching_function
+from flag_engine.segments.utils import get_matching_function
 from flag_engine.utils.hashing import get_hashed_percentage_for_object_ids
 from flag_engine.utils.semver import is_semver
 from flag_engine.utils.types import SupportsStr, get_casting_function
@@ -268,9 +268,8 @@ def get_context_value(
     value = None
     if property.startswith("$."):
         value = _get_context_value_getter(property)(context)
-    elif identity_context := context.get("identity"):
-        if traits := identity_context.get("traits"):
-            value = traits.get(property)
+    else:
+        value = _get_trait_value(context, property)
     return map_any_value_to_context_value(value)
 
 
@@ -357,6 +356,16 @@ def inner(
 }
 
 
+def _get_trait_value(
+    context: EvaluationContext[SegmentMetadataT],
+    trait_key: str,
+) -> ContextValue:
+    if identity_context := context.get("identity"):
+        if traits := identity_context.get("traits"):
+            return traits.get(trait_key)
+    return None
+
+
 @lru_cache
 def _get_context_value_getter(
     property: str,
@@ -373,11 +382,12 @@ def _get_context_value_getter(
     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)}"]',
-        )
+        return partial(_get_trait_value, trait_key=property)
 
     def getter(context: EvaluationContext[SegmentMetadataT]) -> ContextValue:
+        value: object
+        if (value := _get_trait_value(context, property)) is not None:
+            return value
         if typing.TYPE_CHECKING:  # pragma: no cover
             # Ugly hack to satisfy mypy :(
             data = dict(context)

flag_engine/segments/utils.py

@@ -15,10 +15,3 @@ def get_matching_function(
 
 def none(iterable: typing.Iterable[object]) -> bool:
     return not any(iterable)
-
-
-def escape_double_quotes(value: str) -> str:
-    """
-    Escape double quotes in a string for JSONPath compatibility.
-    """
-    return value.replace('"', '\\"')