From 88a852c6f567b5f9c3b17a623fe9b17ffbad8a17 Mon Sep 17 00:00:00 2001 From: dylan Date: Sun, 27 Jul 2025 17:23:40 -0700 Subject: [PATCH] formatting --- example.py | 17 ++++- posthog/client.py | 33 ++++++++- posthog/test/test_client.py | 143 ++++++++++++++++++++++++++++++++++++ posthog/types.py | 3 + 4 files changed, 193 insertions(+), 3 deletions(-) diff --git a/example.py b/example.py index 7311a3cd..eeff5b01 100644 --- a/example.py +++ b/example.py @@ -39,7 +39,7 @@ ) -# Capture an event +# Capture an event with all feature flags posthog.capture( "event", distinct_id="distinct_id", @@ -47,6 +47,21 @@ send_feature_flags=True, ) +# Capture an event with specific feature flags using flag_keys +posthog.capture( + "event-with-specific-flags", + distinct_id="distinct_id", + properties={"property1": "value", "property2": "value"}, + send_feature_flags={ + "only_evaluate_locally": True, + "flag_keys": [ + "beta-feature", + "person-on-events-enabled", + ], # Only evaluate these two flags + "person_properties": {"plan": "premium"}, + }, +) + print(posthog.feature_enabled("beta-feature", "distinct_id")) print( posthog.feature_enabled( diff --git a/posthog/client.py b/posthog/client.py index 65715c46..0e4ce5dc 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -539,6 +539,7 @@ def capture( group_properties=flag_options["group_properties"], disable_geoip=disable_geoip, only_evaluate_locally=True, + flag_keys=flag_options["flag_keys"], ) else: # Default behavior - use remote evaluation @@ -549,6 +550,15 @@ def capture( group_properties=flag_options["group_properties"], disable_geoip=disable_geoip, ) + + # Filter by flag_keys if provided + if flag_options["flag_keys"] is not None: + flag_keys_set = set(flag_options["flag_keys"]) + feature_variants = { + key: value + for key, value in feature_variants.items() + if key in flag_keys_set + } except Exception as e: self.log.exception( f"[FEATURE FLAGS] Unable to get feature variants: {e}" @@ -561,6 +571,7 @@ def capture( groups=(groups or {}), disable_geoip=disable_geoip, only_evaluate_locally=True, + flag_keys=flag_options.get("flag_keys") if flag_options else None, ) for feature, variant in (feature_variants or {}).items(): @@ -588,7 +599,7 @@ def _parse_send_feature_flags(self, send_feature_flags) -> dict: Returns: dict: Normalized options with keys: should_send, only_evaluate_locally, - person_properties, group_properties + person_properties, group_properties, flag_keys Raises: TypeError: If send_feature_flags is not bool or dict @@ -601,6 +612,7 @@ def _parse_send_feature_flags(self, send_feature_flags) -> dict: ), "person_properties": send_feature_flags.get("person_properties"), "group_properties": send_feature_flags.get("group_properties"), + "flag_keys": send_feature_flags.get("flag_keys"), } elif isinstance(send_feature_flags, bool): return { @@ -608,6 +620,7 @@ def _parse_send_feature_flags(self, send_feature_flags) -> dict: "only_evaluate_locally": None, "person_properties": None, "group_properties": None, + "flag_keys": None, } else: raise TypeError( @@ -1667,6 +1680,7 @@ def get_all_flags( group_properties={}, only_evaluate_locally=False, disable_geoip=None, + flag_keys=None, ) -> Optional[dict[str, Union[bool, str]]]: """ Get all feature flags for a user. @@ -1678,6 +1692,7 @@ def get_all_flags( group_properties: A dictionary of group properties. only_evaluate_locally: Whether to only evaluate locally. disable_geoip: Whether to disable GeoIP for this request. + flag_keys: List of specific feature flag keys to evaluate. Examples: ```python @@ -1694,6 +1709,7 @@ def get_all_flags( group_properties=group_properties, only_evaluate_locally=only_evaluate_locally, disable_geoip=disable_geoip, + flag_keys=flag_keys, ) return response["featureFlags"] @@ -1707,6 +1723,7 @@ def get_all_flags_and_payloads( group_properties={}, only_evaluate_locally=False, disable_geoip=None, + flag_keys=None, ) -> FlagsAndPayloads: """ Get all feature flags and their payloads for a user. @@ -1718,6 +1735,7 @@ def get_all_flags_and_payloads( group_properties: A dictionary of group properties. only_evaluate_locally: Whether to only evaluate locally. disable_geoip: Whether to disable GeoIP for this request. + flag_keys: List of specific feature flag keys to evaluate. Examples: ```python @@ -1741,6 +1759,7 @@ def get_all_flags_and_payloads( groups=groups, person_properties=person_properties, group_properties=group_properties, + flag_keys=flag_keys, ) if fallback_to_decide and not only_evaluate_locally: @@ -1768,6 +1787,7 @@ def _get_all_flags_and_payloads_locally( person_properties={}, group_properties={}, warn_on_unknown_groups=False, + flag_keys=None, ) -> tuple[FlagsAndPayloads, bool]: if self.feature_flags is None and self.personal_api_key: self.load_feature_flags() @@ -1777,7 +1797,16 @@ def _get_all_flags_and_payloads_locally( fallback_to_decide = False # If loading in previous line failed if self.feature_flags: - for flag in self.feature_flags: + flags_to_evaluate = self.feature_flags + + # Filter by flag_keys if provided + if flag_keys is not None: + flag_keys_set = set(flag_keys) + flags_to_evaluate = [ + flag for flag in self.feature_flags if flag["key"] in flag_keys_set + ] + + for flag in flags_to_evaluate: try: flags[flag["key"]] = self._compute_flag_locally( flag, diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index d3cd3efc..53a76c9a 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -890,6 +890,143 @@ def test_capture_with_send_feature_flags_options_default_behavior( msg["properties"]["$feature/default-flag"], "default-value" ) + @mock.patch("posthog.client.flags") + def test_capture_with_send_feature_flags_flag_keys_filter(self, patch_flags): + """Test that flag_keys filters the feature flags that are evaluated""" + with mock.patch("posthog.client.batch_post") as mock_post: + client = Client( + FAKE_TEST_API_KEY, + on_error=self.set_fail, + personal_api_key=FAKE_TEST_API_KEY, + sync_mode=True, + ) + + # Set up multiple local flags + client.feature_flags = [ + { + "id": 1, + "key": "flag-one", + "active": True, + "filters": { + "groups": [ + { + "properties": [], + "rollout_percentage": 100, + } + ], + }, + }, + { + "id": 2, + "key": "flag-two", + "active": True, + "filters": { + "groups": [ + { + "properties": [], + "rollout_percentage": 100, + } + ], + }, + }, + { + "id": 3, + "key": "flag-three", + "active": True, + "filters": { + "groups": [ + { + "properties": [], + "rollout_percentage": 100, + } + ], + }, + }, + ] + + # Only evaluate flag-one and flag-three + send_options = { + "only_evaluate_locally": True, + "flag_keys": ["flag-one", "flag-three"], + } + + msg_uuid = client.capture( + "test event", distinct_id="distinct_id", send_feature_flags=send_options + ) + + self.assertIsNotNone(msg_uuid) + self.assertFalse(self.failed) + + # Check the message only includes flag-one and flag-three + mock_post.assert_called_once() + batch_data = mock_post.call_args[1]["batch"] + msg = batch_data[0] + + # Should have flag-one and flag-three, but not flag-two + self.assertEqual(msg["properties"]["$feature/flag-one"], True) + self.assertEqual(msg["properties"]["$feature/flag-three"], True) + self.assertNotIn("$feature/flag-two", msg["properties"]) + + # Active flags should only include flag-one and flag-three + self.assertEqual( + sorted(msg["properties"]["$active_feature_flags"]), + ["flag-one", "flag-three"], + ) + + @mock.patch("posthog.client.flags") + def test_capture_with_send_feature_flags_flag_keys_remote_evaluation( + self, patch_flags + ): + """Test that flag_keys filters remote evaluation results""" + # Mock remote flags response with multiple flags + patch_flags.return_value = { + "featureFlags": { + "remote-flag-one": "value-one", + "remote-flag-two": "value-two", + "remote-flag-three": "value-three", + } + } + + with mock.patch("posthog.client.batch_post") as mock_post: + client = Client( + FAKE_TEST_API_KEY, + on_error=self.set_fail, + sync_mode=True, + ) + + # Only evaluate remote-flag-one and remote-flag-three + send_options = { + "flag_keys": ["remote-flag-one", "remote-flag-three"], + } + + msg_uuid = client.capture( + "test event", distinct_id="distinct_id", send_feature_flags=send_options + ) + + self.assertIsNotNone(msg_uuid) + self.assertFalse(self.failed) + + # Verify flags() was called + patch_flags.assert_called_once() + + # Check the message only includes remote-flag-one and remote-flag-three + mock_post.assert_called_once() + batch_data = mock_post.call_args[1]["batch"] + msg = batch_data[0] + + # Should have remote-flag-one and remote-flag-three, but not remote-flag-two + self.assertEqual(msg["properties"]["$feature/remote-flag-one"], "value-one") + self.assertEqual( + msg["properties"]["$feature/remote-flag-three"], "value-three" + ) + self.assertNotIn("$feature/remote-flag-two", msg["properties"]) + + # Active flags should only include remote-flag-one and remote-flag-three + self.assertEqual( + sorted(msg["properties"]["$active_feature_flags"]), + ["remote-flag-one", "remote-flag-three"], + ) + @mock.patch("posthog.client.flags") def test_capture_exception_with_send_feature_flags_options(self, patch_flags): """Test that capture_exception also supports SendFeatureFlagsOptions""" @@ -2185,6 +2322,7 @@ def test_parse_send_feature_flags_method(self): "only_evaluate_locally": None, "person_properties": None, "group_properties": None, + "flag_keys": None, } self.assertEqual(result, expected) @@ -2195,6 +2333,7 @@ def test_parse_send_feature_flags_method(self): "only_evaluate_locally": None, "person_properties": None, "group_properties": None, + "flag_keys": None, } self.assertEqual(result, expected) @@ -2203,6 +2342,7 @@ def test_parse_send_feature_flags_method(self): "only_evaluate_locally": True, "person_properties": {"plan": "premium"}, "group_properties": {"company": {"type": "enterprise"}}, + "flag_keys": ["beta-feature", "my-flag"], } result = client._parse_send_feature_flags(options) expected = { @@ -2210,6 +2350,7 @@ def test_parse_send_feature_flags_method(self): "only_evaluate_locally": True, "person_properties": {"plan": "premium"}, "group_properties": {"company": {"type": "enterprise"}}, + "flag_keys": ["beta-feature", "my-flag"], } self.assertEqual(result, expected) @@ -2221,6 +2362,7 @@ def test_parse_send_feature_flags_method(self): "only_evaluate_locally": None, "person_properties": {"user_id": "123"}, "group_properties": None, + "flag_keys": None, } self.assertEqual(result, expected) @@ -2231,6 +2373,7 @@ def test_parse_send_feature_flags_method(self): "only_evaluate_locally": None, "person_properties": None, "group_properties": None, + "flag_keys": None, } self.assertEqual(result, expected) diff --git a/posthog/types.py b/posthog/types.py index 3cc505c0..9d691ba9 100644 --- a/posthog/types.py +++ b/posthog/types.py @@ -20,11 +20,14 @@ class SendFeatureFlagsOptions(TypedDict, total=False): These properties will be merged with any existing person properties. group_properties: Group properties to use for feature flag evaluation specific to this event. Format: { group_type_name: { group_properties } } + flag_keys: List of specific feature flag keys to evaluate. Only these flags will be evaluated + and included in the event. If not provided, all flags will be evaluated. """ only_evaluate_locally: Optional[bool] person_properties: Optional[dict[str, Any]] group_properties: Optional[dict[str, dict[str, Any]]] + flag_keys: Optional[list[str]] @dataclass(frozen=True)