feat: Support $.flags context values in segment conditions

Environment flags are now evaluated with segments whose conditions all read
`$.environment` or `$.flags`, as `GET /api/v1/flags/` does since
Flagsmith/flagsmith#8576. Segments reading traits or identity context are
left out.

When segment conditions read `$.flags`, identity flags are evaluated at
once against the full context, as a context trimmed to one feature lacks
the flags it depends on. Otherwise, they are still evaluated lazily.

Closes #245

Author: khvn26

flagsmith/flagsmith.py

@@ -513,6 +513,12 @@ def _get_identity_flags_from_document(
             identifier=identifier,
             traits=traits,
         )
+        if local_evaluation.reads.flags:
+            return Flags.from_evaluation_result(
+                evaluation_result=engine.get_evaluation_result(context=context),
+                analytics_processor=self._analytics_processor,
+                default_flag_handler=self.default_flag_handler,
+            )
         # Lazy: defer per-feature evaluation until the caller actually reads
         # a flag. Hot for callers that only read one or a few flags out of a
         # large environment.

flagsmith/local_evaluation.py

@@ -5,7 +5,7 @@
 from functools import cached_property
 
 from flag_engine import engine
-from flag_engine.context.types import SegmentContext
+from flag_engine.context.types import SegmentContext, SegmentRule
 
 from flagsmith.api.types import EnvironmentModel
 from flagsmith.mappers import (
@@ -15,6 +15,7 @@
 )
 from flagsmith.models import (
     Flag,
+    SegmentContexts,
     SegmentOverridesIndex,
     build_segment_overrides_index,
 )
@@ -25,6 +26,35 @@
     TraitMapping,
 )
 
+_FLAGS_PROPERTY_PREFIXES = ("$.flags.", "$.flags[")
+_IDENTITY_FREE_PROPERTY_PREFIXES = (
+    "$.environment.",
+    "$.environment[",
+    *_FLAGS_PROPERTY_PREFIXES,
+)
+
+
+class _ContextReads(typing.NamedTuple):
+    """The context values segment rules read."""
+
+    flags: bool
+    identity: bool
+
+
+def _get_context_reads(rules: typing.Iterable[SegmentRule]) -> _ContextReads:
+    reads_flags = reads_identity = False
+    for rule in rules:
+        for condition in rule.get("conditions") or ():
+            condition_property = condition.get("property") or ""
+            reads_flags |= condition_property.startswith(_FLAGS_PROPERTY_PREFIXES)
+            reads_identity |= not condition_property.startswith(
+                _IDENTITY_FREE_PROPERTY_PREFIXES
+            )
+        nested_reads = _get_context_reads(rule.get("rules") or ())
+        reads_flags |= nested_reads.flags
+        reads_identity |= nested_reads.identity
+    return _ContextReads(flags=reads_flags, identity=reads_identity)
+
 
 class LocalEvaluation:
     """Per-environment-document state for local evaluation."""
@@ -37,6 +67,15 @@ def __init__(
         ] = None,
     ) -> None:
         self.context = context
+        self._identity_free_segments: SegmentContexts = {}
+        reads_flags = reads_identity = False
+        for segment_key, segment_context in (context.get("segments") or {}).items():
+            reads = _get_context_reads(segment_context["rules"])
+            reads_flags |= reads.flags
+            reads_identity |= reads.identity
+            if not reads.identity:
+                self._identity_free_segments[segment_key] = segment_context
+        self.reads = _ContextReads(flags=reads_flags, identity=reads_identity)
         self.overrides_index: SegmentOverridesIndex = build_segment_overrides_index(
             context
         )
@@ -64,10 +103,8 @@ def from_environment_document(
     @cached_property
     def environment_flags(self) -> typing.Dict[str, Flag]:
         """Environment flags, which depend on the environment document only."""
-        context_without_segments = self.context.copy()
-        context_without_segments.pop("segments", None)
         evaluation_result = engine.get_evaluation_result(
-            context=context_without_segments,
+            context={**self.context, "segments": self._identity_free_segments},
         )
         return {
             feature_name: Flag.from_evaluation_result(flag_result)

tests/test_flagsmith.py

@@ -18,6 +18,7 @@
     FlagsmithFeatureDoesNotExistError,
 )
 from flagsmith.local_evaluation import LocalEvaluation
+from flagsmith.mappers import map_environment_document_to_context
 from flagsmith.models import DefaultFlag, ExperimentMetadata, Flag, Flags
 from flagsmith.offline_handlers import OfflineHandler
 from flagsmith.types import SDKEvaluationContext
@@ -101,40 +102,72 @@ def test_get_environment_flags_uses_local_environment_when_available(
     assert all_flags[0].value == "some-value"
 
 
-def test_get_environment_flags_omits_segments_from_evaluation_context(
+def _condition(property_: str) -> typing.Any:
+    return {"operator": "EQUAL", "property_": property_, "value": "value"}
+
+
+@pytest.mark.parametrize(
+    "rules, expected_segment_names",
+    [
+        (
+            [{"type": "ALL", "conditions": [_condition("$.environment.name")]}],
+            ["Test segment"],
+        ),
+        (
+            [
+                {
+                    "type": "ALL",
+                    "rules": [
+                        {
+                            "type": "ANY",
+                            "conditions": [
+                                _condition("$.environment['name']"),
+                                _condition("$.flags.some_feature.enabled"),
+                            ],
+                        }
+                    ],
+                }
+            ],
+            ["Test segment"],
+        ),
+        ([{"type": "ALL", "conditions": [_condition("trait")]}], []),
+        (
+            [
+                {
+                    "type": "ALL",
+                    "conditions": [_condition("$.environment.name")],
+                    "rules": [{"type": "ALL", "conditions": [_condition("trait")]}],
+                }
+            ],
+            [],
+        ),
+    ],
+)
+def test_get_environment_flags__local_evaluation__includes_identity_free_segments_only(
     mocker: MockerFixture,
-    local_eval_flagsmith: Flagsmith,
-    evaluation_context: SDKEvaluationContext,
+    flagsmith: Flagsmith,
+    environment: EnvironmentModel,
+    rules: typing.Any,
+    expected_segment_names: typing.List[str],
 ) -> None:
     # Given
-    mock_get_evaluation_result = mocker.patch(
-        "flagsmith.flagsmith.engine.get_evaluation_result",
-        autospec=True,
+    # the only segment has the given rules, plus identity overrides
+    environment["project"]["segments"][0]["rules"] = rules
+    flagsmith._local_evaluation = LocalEvaluation(
+        map_environment_document_to_context(environment)
     )
-
-    expected_evaluation_result = {
-        "flags": {
-            "some_feature": {
-                "name": "some_feature",
-                "enabled": True,
-                "value": "some-feature-state-value",
-                "reason": "DEFAULT",
-                "metadata": {"id": 1},
-            }
-        },
-        "segments": [],
-    }
-
-    mock_get_evaluation_result.return_value = expected_evaluation_result
+    flagsmith.enable_local_evaluation = True
+    spy = mocker.spy(engine, "get_evaluation_result")
 
     # When
-    local_eval_flagsmith.get_environment_flags()
+    flagsmith.get_environment_flags()
 
     # Then
-    # Verify segments are not present in the context passed to the engine
-    context_without_segments = evaluation_context.copy()
-    context_without_segments.pop("segments", None)
-    mock_get_evaluation_result.assert_called_once_with(context=context_without_segments)
+    # only segments reading `$.environment` or `$.flags` alone are evaluated,
+    # as `GET /api/v1/flags/` does
+    context = spy.call_args.kwargs["context"]
+    segments = (context.get("segments") or {}).values()
+    assert [segment["name"] for segment in segments] == expected_segment_names
 
 
 def test_get_environment_flags__local_evaluation__evaluates_once_per_environment(
@@ -316,6 +349,40 @@ def test_get_identity_flags_includes_segments_in_evaluation_context(
     assert "segments" in context
 
 
+def test_get_identity_flags__flags_condition__evaluates_full_context(
+    mocker: MockerFixture,
+    flagsmith: Flagsmith,
+    environment: EnvironmentModel,
+) -> None:
+    # Given
+    # a segment condition reading another flag
+    environment["project"]["segments"][0]["rules"] = [
+        {
+            "type": "ALL",
+            "rules": [],
+            "conditions": [_condition("$.flags.some_feature.enabled")],
+        }
+    ]
+    flagsmith._local_evaluation = LocalEvaluation(
+        map_environment_document_to_context(environment)
+    )
+    flagsmith.enable_local_evaluation = True
+    spy = mocker.spy(engine, "get_evaluation_result")
+
+    # When
+    flagsmith.get_identity_flags("identifier").get_flag("some_feature")
+
+    # Then
+    # flags are evaluated at once against the full context, segments included,
+    # rather than per flag against a trimmed one
+    spy.assert_called_once()
+    context = spy.call_args.kwargs["context"]
+    assert set(context["features"]) == {"some_feature"}
+    assert [segment["name"] for segment in (context["segments"] or {}).values()] == [
+        "Test segment"
+    ]
+
+
 def test_get_identity_flags__resolves_one_flag_at_a_time(
     local_eval_flagsmith: Flagsmith,
     mocker: MockerFixture,

tests/test_local_evaluation.py

@@ -1,3 +1,5 @@
+import copy
+
 import pytest
 from flag_engine import engine
 from pytest_mock import MockerFixture
@@ -141,3 +143,32 @@ def test_local_evaluation__environment_flags__evaluated_once(
     assert first is second
     assert first["some_feature"].value == "some-value"
     spy.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "condition_property, expected_reads_flags",
+    [
+        ("$.flags.some_feature.enabled", True),
+        ("$.flags['some feature'].value", True),
+        ("$.flagship", False),
+        ("flags", False),
+    ],
+)
+def test_local_evaluation__reads__flags__returns_expected(
+    environment: EnvironmentModel,
+    condition_property: str,
+    expected_reads_flags: bool,
+) -> None:
+    # Given
+    # a condition in a nested rule of the only segment
+    environment = copy.deepcopy(environment)
+    condition = environment["project"]["segments"][0]["rules"][0]["rules"][0][
+        "conditions"
+    ][0]
+    condition["property_"] = condition_property
+
+    # When
+    local_evaluation = LocalEvaluation.from_environment_document(environment)
+
+    # Then
+    assert local_evaluation.reads.flags is expected_reads_flags