feat(v7): JSONPath support

Author: khvn26

flag_engine/segments/evaluator.py

@@ -4,8 +4,9 @@
 import typing
 import warnings
 from contextlib import suppress
-from functools import partial, wraps
+from functools import lru_cache, wraps
 
+import jsonpath_rfc9535
 import semver
 
 from flag_engine.context.mappers import map_environment_identity_to_context
@@ -23,7 +24,7 @@
 from flag_engine.segments import constants
 from flag_engine.segments.models import SegmentModel
 from flag_engine.segments.types import ConditionOperator
-from flag_engine.segments.utils import get_matching_function
+from flag_engine.segments.utils import escape_double_quotes, 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
@@ -256,26 +257,11 @@ def context_matches_condition(
     )
 
 
-def _get_trait(context: EvaluationContext, trait_key: str) -> ContextValue:
-    return (
-        identity_context["traits"][trait_key]
-        if (identity_context := context["identity"])
-        else None
-    )
-
-
 def get_context_value(
     context: EvaluationContext,
     property: str,
 ) -> ContextValue:
-    getter = CONTEXT_VALUE_GETTERS_BY_PROPERTY.get(property) or partial(
-        _get_trait,
-        trait_key=property,
-    )
-    try:
-        return getter(context)
-    except KeyError:
-        return None
+    return _get_context_value_getter(property)(context)
 
 
 def _matches_context_value(
@@ -385,8 +371,36 @@ def inner(
 }
 
 
-CONTEXT_VALUE_GETTERS_BY_PROPERTY = {
-    "$.identity.identifier": lambda context: context["identity"]["identifier"],
-    "$.identity.key": lambda context: context["identity"]["key"],
-    "$.environment.name": lambda context: context["environment"]["name"],
-}
+@lru_cache
+def _get_context_value_getter(
+    property: str,
+) -> typing.Callable[[EvaluationContext], ContextValue]:
+    """
+    Get a function to retrieve a context value based on property value,
+    assumed to be either a JSONPath string or a trait key.
+
+    :param property: The property to retrieve the value for.
+    :return: A function that takes an EvaluationContext and returns the value.
+    """
+    try:
+        compiled_query = jsonpath_rfc9535.compile(property)
+    except jsonpath_rfc9535.JSONPathSyntaxError:
+        compiled_query = jsonpath_rfc9535.compile(
+            f'$.identity.traits["{escape_double_quotes(property)}"]',
+        )
+
+    def getter(context: EvaluationContext) -> ContextValue:
+        try:
+            if result := compiled_query.find_one(context):
+                return result.value
+            return None
+        except jsonpath_rfc9535.JSONPathError:  # pragma: no cover
+            # This is supposed to be unreachable, but if it happens,
+            # we log a warning and return None.
+            warnings.warn(
+                f"Failed to evaluate JSONPath query '{property}' in context: {context}",
+                RuntimeWarning,
+            )
+            return None
+
+    return getter

flag_engine/segments/utils.py

@@ -15,3 +15,10 @@ 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('"', '\\"')

requirements.in

@@ -1,5 +1,6 @@
 annotated-types
-semver
+jsonpath-rfc9535
 pydantic
 pydantic-collections
+semver
 typing_extensions

requirements.txt

@@ -1,26 +1,31 @@
-#
-# This file is autogenerated by pip-compile with Python 3.11
-# by the following command:
-#
-#    pip-compile
-#
-annotated-types==0.5.0
+# This file was autogenerated by uv via the following command:
+#    uv pip compile requirements.in --constraints requirements.txt
+annotated-types==0.7.0
     # via
     #   -r requirements.in
     #   pydantic
-pydantic==2.4.0
+iregexp-check==0.1.4
+    # via jsonpath-rfc9535
+jsonpath-rfc9535==0.1.5
+    # via -r requirements.in
+pydantic==2.11.7
     # via
     #   -r requirements.in
     #   pydantic-collections
-pydantic-collections==0.5.1
+pydantic-collections==0.6.0
     # via -r requirements.in
-pydantic-core==2.10.0
+pydantic-core==2.33.2
     # via pydantic
-semver==3.0.1
+regex==2025.7.34
+    # via jsonpath-rfc9535
+semver==3.0.4
     # via -r requirements.in
-typing-extensions==4.8.0
+typing-extensions==4.14.1
     # via
     #   -r requirements.in
     #   pydantic
     #   pydantic-collections
     #   pydantic-core
+    #   typing-inspection
+typing-inspection==0.4.1
+    # via pydantic

setup.py

@@ -13,6 +13,7 @@
     long_description=open("README.md").read(),
     long_description_content_type="text/markdown",
     install_requires=[
+        "jsonpath-rfc9535>=0.1.5,<1",
         "pydantic>=2.3.0,<3",
         "pydantic-collections>=0.5.1,<1",
         "semver>=3.0.1",