feat(v7): JSONPath support (#243)

Author: khvn26

.github/workflows/pull-request.yml

@@ -30,18 +30,17 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
-
       - name: Install Dependencies
-        if: ${{ matrix.python-version == '3.8' }}
+        if: ${{ matrix.python-version != '3.8' }}
         run: |
           python -m pip install --upgrade pip
-          pip install -r requirements.txt -r requirements-dev-3.8.txt
+          pip install -r requirements.txt -r requirements-dev.txt
 
       - name: Install Dependencies
-        if: ${{ matrix.python-version != '3.8' }}
+        if: ${{ matrix.python-version == '3.8' }}
         run: |
           python -m pip install --upgrade pip
-          pip install -r requirements.txt -r requirements-dev.txt
+          pip install -r requirements.txt -r requirements-dev-3.8.txt
 
       - name: Check Formatting
         run: black --check .

flag_engine/identities/traits/types.py

@@ -11,6 +11,14 @@
 _UnconstrainedContextValue = Union[None, int, float, bool, str]
 
 
+def is_trait_value(value: Any) -> TypeGuard[_UnconstrainedContextValue]:
+    """
+    Check if the value is a valid trait value type.
+    This function is used to determine if a value can be treated as a trait value.
+    """
+    return isinstance(value, get_args(_UnconstrainedContextValue))
+
+
 def map_any_value_to_trait_value(value: Any) -> _UnconstrainedContextValue:
     """
     Try to coerce a value of arbitrary type to a trait value type.

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
@@ -18,12 +19,12 @@
 )
 from flag_engine.environments.models import EnvironmentModel
 from flag_engine.identities.models import IdentityModel
-from flag_engine.identities.traits.types import ContextValue
+from flag_engine.identities.traits.types import ContextValue, is_trait_value
 from flag_engine.result.types import EvaluationResult, FlagResult, SegmentResult
 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,16 @@ 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
+    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,8 +376,44 @@ 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:
+        # 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)}"]',
+        )
+
+    def getter(context: EvaluationContext) -> ContextValue:
+        if typing.TYPE_CHECKING:  # pragma: no cover
+            # Ugly hack to satisfy mypy :(
+            data = dict(context)
+        else:
+            data = context
+        try:
+            if result := compiled_query.find_one(data):
+                if is_trait_value(value := result.value):
+                    return 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-dev-3.8.txt

@@ -2,7 +2,7 @@
 #    uv pip compile requirements-dev.in --constraints requirements.txt -o requirements-dev.txt --python-version 3.8
 absolufy-imports==0.3.1
     # via -r requirements-dev.in
-annotated-types==0.5.0
+annotated-types==0.7.0
     # via
     #   -c requirements.txt
     #   pydantic
@@ -74,11 +74,11 @@ pycodestyle==2.9.1
     # via flake8
 pycparser==2.22
     # via cffi
-pydantic==2.4.0
+pydantic==2.10.6
     # via
     #   -c requirements.txt
     #   datamodel-code-generator
-pydantic-core==2.10.0
+pydantic-core==2.27.2
     # via
     #   -c requirements.txt
     #   pydantic
@@ -118,7 +118,7 @@ tomli==2.2.1
     #   pytest
 types-setuptools==75.8.0.20250110
     # via -r requirements-dev.in
-typing-extensions==4.8.0
+typing-extensions==4.13.2
     # via
     #   -c requirements.txt
     #   annotated-types

requirements-dev.txt

@@ -2,7 +2,7 @@
 #    uv pip compile requirements-dev.in --constraints requirements.txt -o requirements-dev.txt --python-version 3.9
 absolufy-imports==0.3.1
     # via -r requirements-dev.in
-annotated-types==0.5.0
+annotated-types==0.7.0
     # via
     #   -c requirements.txt
     #   pydantic
@@ -85,11 +85,11 @@ pycodestyle==2.14.0
     # via flake8
 pycparser==2.22
     # via cffi
-pydantic==2.4.0
+pydantic==2.10.6
     # via
     #   -c requirements.txt
     #   datamodel-code-generator
-pydantic-core==2.10.0
+pydantic-core==2.27.2
     # via
     #   -c requirements.txt
     #   pydantic
@@ -133,11 +133,11 @@ tomli==2.2.1
     #   mypy
     #   pip-tools
     #   pytest
-typeguard==4.2.0
+typeguard==4.4.2
     # via inflect
 types-setuptools==80.9.0.20250809
     # via -r requirements-dev.in
-typing-extensions==4.8.0
+typing-extensions==4.13.2
     # via
     #   -c requirements.txt
     #   black

requirements.in

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

requirements.txt

@@ -1,26 +1,29 @@
-#
-# 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 --python-version 3.8
+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.6
+    # via -r requirements.in
+pydantic==2.10.6
     # 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.27.2
     # via pydantic
-semver==3.0.1
+regex==2024.11.6
+    # via jsonpath-rfc9535
+semver==3.0.4
     # via -r requirements.in
-typing-extensions==4.8.0
+typing-extensions==4.13.2
     # via
     #   -r requirements.in
+    #   annotated-types
     #   pydantic
     #   pydantic-collections
     #   pydantic-core

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",

tests/unit/segments/test_segments_evaluator.py

@@ -21,6 +21,7 @@
     _matches_context_value,
     context_matches_condition,
     get_context_segments,
+    get_context_value,
     get_evaluation_result,
     get_flag_result_from_feature_context,
     get_identity_segments,
@@ -1028,3 +1029,38 @@ def test_get_flag_result_from_feature_context__call_return_expected(
             expected_key,
         ]
     )
+
+
+@pytest.mark.parametrize(
+    "property",
+    [
+        pytest.param("$.identity", id="jsonpath returns an object"),
+        'trait key "with quotes"',
+        "$.leads.to.nowhere",
+    ],
+)
+def test_get_context_value__invalid_jsonpath__returns_expected(
+    context: EvaluationContext,
+    property: str,
+) -> None:
+    # Given & When
+    result = get_context_value(context, property)
+
+    # 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