From b6b81cfb6957b9fb5bbd0731852a8f4abba9cc64 Mon Sep 17 00:00:00 2001 From: OpenHands Agent Date: Wed, 21 Oct 2026 16:57:20 +0000 Subject: [PATCH] feat: Support registering global traits at client initialization Add a `default_traits` argument to `Flagsmith.__init__` so traits that never change for the life of a client (app name, developer branch, and so on) can be registered once instead of being supplied on every call. Default traits are merged into `get_identity_flags`, `get_identity_segments`, `get_experiment_flag`, `track_event` and `track_exposure_event`. Traits passed to an individual call take precedence on key conflicts, and default traits accept the same `{"value": ..., "transient": True}` envelope as per-call traits. Closes #237 Co-Authored-By: Claude Opus 5 (1M context) --- flagsmith/flagsmith.py | 31 ++++- tests/test_flagsmith.py | 242 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 268 insertions(+), 5 deletions(-) diff --git a/flagsmith/flagsmith.py b/flagsmith/flagsmith.py index ee028a2..c1a3b6e 100644 --- a/flagsmith/flagsmith.py +++ b/flagsmith/flagsmith.py @@ -78,6 +78,7 @@ def __init__( offline_handler: typing.Optional[OfflineHandler] = None, enable_realtime_updates: bool = False, application_metadata: typing.Optional[ApplicationMetadata] = None, + default_traits: typing.Optional[TraitMapping] = None, ): """ :param environment_key: The environment key obtained from Flagsmith interface. @@ -117,6 +118,12 @@ def __init__( default_flag_handler if offline_mode is not set and using remote evaluation. :param enable_realtime_updates: Use real-time functionality via SSE as opposed to polling the API :param application_metadata: Optional metadata about the client application. + :param default_traits: traits that apply to every identity evaluation and + event for the lifetime of this client, e.g. + `{"app_name": "my-app", "branch": "main"}`. Traits passed to an + individual call take precedence over a default trait of the same key. + Envelope traits you don't want persisted in a dictionary with + `"transient"` and `"value"` keys, as for per-call traits. """ self.offline_mode = offline_mode @@ -125,6 +132,7 @@ def __init__( self.offline_handler = offline_handler self.default_flag_handler = default_flag_handler self.enable_realtime_updates = enable_realtime_updates + self.default_traits: TraitMapping = default_traits or {} self._analytics_processor: typing.Optional[AnalyticsProcessor] = None self._event_processor: typing.Optional[EventProcessor] = None self._local_evaluation: typing.Optional[LocalEvaluation] = None @@ -207,6 +215,15 @@ def __init__( def _ensure_trailing_slash(url: str) -> str: return url if url.endswith("/") else f"{url}/" + def _merge_traits(self, traits: typing.Optional[TraitMapping]) -> TraitMapping: + """Combine the client's default traits with traits given for a single + call. Per-call traits win on key conflicts.""" + if not self.default_traits: + return traits or {} + if not traits: + return self.default_traits + return {**self.default_traits, **traits} + def _initialise_analytics( self, environment_key: str, @@ -305,10 +322,11 @@ def get_identity_flags( Flagsmith, e.g. `{"num_orders": 10}`. Envelope traits you don't want persisted in a dictionary with `"transient"` and `"value"` keys, e.g. `{"num_orders": 10, "color": {"value": "pink", "transient": True}}`. + Merged over the client's `default_traits`, if any. :param transient: if `True`, the identity won't get persisted :return: Flags object holding all the flags for the given identity. """ - traits = traits or {} + traits = self._merge_traits(traits) if ( self.offline_mode or self.enable_local_evaluation ) and self._local_evaluation: @@ -330,7 +348,8 @@ def get_identity_segments( :param identifier: a unique identifier for the identity in the current environment, e.g. email address, username, uuid :param traits: a dictionary of traits to add / update on the identity in - Flagsmith, e.g. {"num_orders": 10} + Flagsmith, e.g. {"num_orders": 10}. Merged over the client's + `default_traits`, if any. :return: list of Segment objects that the identity is part of. """ if not (local_evaluation := self._local_evaluation): @@ -341,7 +360,7 @@ def get_identity_segments( context = map_context_and_identity_data_to_context( context=local_evaluation.context, identifier=identifier, - traits=traits, + traits=self._merge_traits(traits), ) evaluation_result = engine.get_evaluation_result( @@ -366,6 +385,8 @@ def get_experiment_flag( other outcome is logged and skipped to keep experimentation data clean. A `DefaultFlag` served via the `default_flag_handler` counts as the feature not existing. + + `traits` are merged over the client's `default_traits`, if any. """ if not self._event_processor: raise ValueError("Events must be enabled to use experiment flags.") @@ -423,7 +444,7 @@ def track_event( event=event, identifier=identifier, value=value, - traits=resolve_trait_values(traits), + traits=resolve_trait_values(self._merge_traits(traits)), metadata=metadata, ) @@ -449,7 +470,7 @@ def track_exposure_event( feature_name=feature_name, identifier=identifier, value=value, - traits=resolve_trait_values(traits), + traits=resolve_trait_values(self._merge_traits(traits)), metadata=metadata, ) diff --git a/tests/test_flagsmith.py b/tests/test_flagsmith.py index b3cbfdd..a147ca5 100644 --- a/tests/test_flagsmith.py +++ b/tests/test_flagsmith.py @@ -1455,3 +1455,245 @@ def test_get_experiment_flag__disabled_feature__skips_exposure( "Not sending $flag_exposure for feature checkout_v2: feature is disabled." in caplog.messages ) + + +def test_default_traits__not_set__empty(api_key: str) -> None: + # Given / When + flagsmith = Flagsmith(environment_key=api_key) + + # Then + assert flagsmith.default_traits == {} + + +@responses.activate() +def test_default_traits__get_identity_flags__sent_to_api( + api_key: str, identities_json: str +) -> None: + # Given + flagsmith = Flagsmith( + environment_key=api_key, + default_traits={ + "app_name": "my-app", + "branch": {"value": "main", "transient": True}, + }, + ) + responses.add( + method="POST", + url=flagsmith.identities_url, + body=identities_json, + match=[ + matchers.json_params_matcher( + { + "identifier": "identifier", + "traits": [ + {"trait_key": "app_name", "trait_value": "my-app"}, + { + "trait_key": "branch", + "trait_value": "main", + "transient": True, + }, + {"trait_key": "num_orders", "trait_value": 10}, + ], + } + ) + ], + ) + + # When & Then + flagsmith.get_identity_flags("identifier", traits={"num_orders": 10}) + + +@responses.activate() +def test_default_traits__get_identity_flags__per_call_traits_take_precedence( + api_key: str, identities_json: str +) -> None: + # Given + flagsmith = Flagsmith( + environment_key=api_key, + default_traits={"app_name": "my-app"}, + ) + responses.add( + method="POST", + url=flagsmith.identities_url, + body=identities_json, + match=[ + matchers.json_params_matcher( + { + "identifier": "identifier", + "traits": [ + {"trait_key": "app_name", "trait_value": "another-app"}, + ], + } + ) + ], + ) + + # When & Then + flagsmith.get_identity_flags("identifier", traits={"app_name": "another-app"}) + + +def test_default_traits__local_evaluation__included_in_context( + requests_session_response_ok: None, server_api_key: str, mocker: MockerFixture +) -> None: + # Given + flagsmith = Flagsmith( + environment_key=server_api_key, + enable_local_evaluation=True, + environment_refresh_interval_seconds=0.1, + default_traits={ + "app_name": "my-app", + "branch": {"value": "main", "transient": True}, + }, + ) + spy = mocker.spy(engine, "get_evaluation_result") + + # When + flagsmith.get_identity_flags("identifier", traits={"num_orders": 10}).all_flags() + + # Then + context = spy.call_args.kwargs.get("context") or spy.call_args.args[0] + assert context["identity"]["traits"] == { + "app_name": "my-app", + "branch": "main", + "num_orders": 10, + } + + +def test_default_traits__get_identity_segments__matched_without_per_call_traits( + requests_session_response_ok: None, server_api_key: str +) -> None: + # Given + flagsmith = Flagsmith( + environment_key=server_api_key, + enable_local_evaluation=True, + environment_refresh_interval_seconds=0.1, + default_traits={"foo": "bar"}, # obtained from data/environment.json + ) + + # When + segments = flagsmith.get_identity_segments("identifier") + + # Then + assert [segment.name for segment in segments] == ["Test segment"] + + +def test_default_traits__get_identity_segments__per_call_traits_take_precedence( + requests_session_response_ok: None, server_api_key: str +) -> None: + # Given + flagsmith = Flagsmith( + environment_key=server_api_key, + enable_local_evaluation=True, + environment_refresh_interval_seconds=0.1, + default_traits={"foo": "not-bar"}, + ) + + # When + segments = flagsmith.get_identity_segments("identifier", traits={"foo": "bar"}) + + # Then + assert [segment.name for segment in segments] == ["Test segment"] + + +def test_default_traits__track_event__merged_with_call_traits( + mocker: MockerFixture, api_key: str +) -> None: + # Given + flagsmith = Flagsmith( + environment_key=api_key, + enable_events=True, + event_processor_config=EventProcessorConfig(events_api_url="http://test/"), + default_traits={ + "app_name": "my-app", + "plan": {"value": "free", "transient": True}, + }, + ) + mock_track = mocker.patch.object(flagsmith._event_processor, "track_event") + + # When + flagsmith.track_event("purchase", identifier="user1", traits={"plan": "premium"}) + + # Then + mock_track.assert_called_once_with( + event="purchase", + identifier="user1", + value=None, + traits={"app_name": "my-app", "plan": "premium"}, + metadata=None, + ) + + +def test_default_traits__track_exposure_event__merged_with_call_traits( + mocker: MockerFixture, api_key: str +) -> None: + # Given + flagsmith = Flagsmith( + environment_key=api_key, + enable_events=True, + event_processor_config=EventProcessorConfig(events_api_url="http://test/"), + default_traits={"app_name": "my-app"}, + ) + mock_track = mocker.patch.object(flagsmith._event_processor, "track_exposure_event") + + # When + flagsmith.track_exposure_event( + "checkout_v2", + identifier="user1", + value="variant_b", + traits={"plan": "premium"}, + ) + + # Then + mock_track.assert_called_once_with( + feature_name="checkout_v2", + identifier="user1", + value="variant_b", + traits={"app_name": "my-app", "plan": "premium"}, + metadata=None, + ) + + +@responses.activate() +def test_default_traits__get_experiment_flag__merged_once_per_call( + mocker: MockerFixture, api_key: str, identities_json: str +) -> None: + # Given + flagsmith = Flagsmith( + environment_key=api_key, + enable_events=True, + event_processor_config=EventProcessorConfig(events_api_url="http://test/"), + default_traits={"app_name": "my-app"}, + ) + mock_track = mocker.patch.object(flagsmith._event_processor, "track_exposure_event") + responses.add( + method="POST", + url=flagsmith.identities_url, + body=identities_json, + match=[ + matchers.json_params_matcher( + { + "identifier": "user1", + "traits": [ + {"trait_key": "app_name", "trait_value": "my-app"}, + {"trait_key": "plan", "trait_value": "premium"}, + ], + } + ) + ], + ) + + # When + flagsmith.get_experiment_flag( + feature_name="some_feature", + identifier="user1", + traits={"plan": "premium"}, + ) + + # Then + mock_track.assert_called_once_with( + feature_name="some_feature", + identifier="user1", + value="treatment", + traits={"app_name": "my-app", "plan": "premium"}, + metadata={"experiment_id": 42}, + )