fix typing, add type guard

Author: khvn26

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

@@ -6,7 +6,7 @@
 from contextlib import suppress
 from functools import lru_cache, wraps
 
-import jsonpath_rfc9535  # type: ignore[import-untyped]
+import jsonpath_rfc9535  # type: ignore[import]
 import semver
 
 from flag_engine.context.mappers import map_environment_identity_to_context
@@ -19,7 +19,7 @@
 )
 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
@@ -392,7 +392,8 @@ def _get_context_value_getter(
     def getter(context: EvaluationContext) -> ContextValue:
         try:
             if result := compiled_query.find_one(context):
-                return result.value
+                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,

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,22 @@ 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