From 9a00bf287f235343606851ac3051fbb8e10b3615 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Fri, 19 Jun 2026 12:33:52 +0200 Subject: [PATCH 1/4] fix: Add context helpers to Client --- .changeset/fuzzy-pandas-scope.md | 5 ++ posthog/client.py | 113 ++++++++++++++++++++++++++++- posthog/test/test_client.py | 81 +++++++++++++++++++++ references/public_api_snapshot.txt | 6 ++ 4 files changed, 201 insertions(+), 4 deletions(-) create mode 100644 .changeset/fuzzy-pandas-scope.md diff --git a/.changeset/fuzzy-pandas-scope.md b/.changeset/fuzzy-pandas-scope.md new file mode 100644 index 00000000..25f9053d --- /dev/null +++ b/.changeset/fuzzy-pandas-scope.md @@ -0,0 +1,5 @@ +--- +'pypi/posthog': patch +--- + +Add context helper methods to custom PostHog client instances. diff --git a/posthog/client.py b/posthog/client.py index 800cc577..6ceb7850 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -8,7 +8,7 @@ import warnings import weakref from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Optional, TypeVar, Union, cast from uuid import uuid4 from typing_extensions import Unpack @@ -24,7 +24,12 @@ get_context_device_id, get_context_distinct_id, get_context_session_id, + get_tags as _context_get_tags, + identify_context as _context_identify_context, new_context, + set_context_device_id as _context_set_context_device_id, + set_context_session as _context_set_context_session, + tag as _context_tag, ) from posthog.exception_capture import ExceptionCapture from posthog.exception_utils import ( @@ -99,6 +104,8 @@ MAX_DICT_SIZE = 50_000 +_F = TypeVar("_F", bound=Callable[..., Any]) + def get_identity_state(passed) -> tuple[str, bool]: """Returns the distinct id to use, and whether this is a personless event or not""" @@ -466,9 +473,9 @@ def new_context(self, fresh=False, capture_exceptions=True): Examples: ```python - with posthog.new_context(): - identify_context('') - posthog.capture('event_name') + with client.new_context(): + client.identify_context('') + client.capture('event_name') ``` Category: @@ -478,6 +485,104 @@ def new_context(self, fresh=False, capture_exceptions=True): fresh=fresh, capture_exceptions=capture_exceptions, client=self ) + def scoped(self, fresh=False, capture_exceptions=True): + """ + Decorator that creates a new context for the wrapped function using this client. + + Args: + fresh: Whether to create a fresh context that doesn't inherit from parent. + capture_exceptions: Whether to automatically capture exceptions in this context. + + Category: + Contexts + """ + + def decorator(func: _F) -> _F: + from functools import wraps + + if inspect.iscoroutinefunction(func): + + @wraps(func) + async def async_wrapper(*args, **kwargs): + with self.new_context( + fresh=fresh, capture_exceptions=capture_exceptions + ): + return await func(*args, **kwargs) + + return cast(_F, async_wrapper) + + @wraps(func) + def wrapper(*args, **kwargs): + with self.new_context( + fresh=fresh, capture_exceptions=capture_exceptions + ): + return func(*args, **kwargs) + + return cast(_F, wrapper) + + return decorator + + def tag(self, name: str, value: Any) -> None: + """ + Add a tag to the current context. + + Args: + name: The tag key. + value: The tag value. + + Category: + Contexts + """ + _context_tag(name, value) + + def get_tags(self) -> Dict[str, Any]: + """ + Get all tags from the current context. + + Returns: + Dict of all tags in the current context. + + Category: + Contexts + """ + return _context_get_tags() + + def identify_context(self, distinct_id: str) -> None: + """ + Identify the current context with a distinct ID. + + Args: + distinct_id: The distinct ID to associate with the current context and its children. + + Category: + Contexts + """ + _context_identify_context(distinct_id) + + def set_context_session(self, session_id: str) -> None: + """ + Set the session ID for the current context. + + Args: + session_id: The session ID to associate with the current context and its children. + + Category: + Contexts + """ + _context_set_context_session(session_id) + + def set_context_device_id(self, device_id: str) -> None: + """ + Set the device ID for the current context. + + Args: + device_id: The device ID to associate with the current context and its children. + + Category: + Contexts + """ + _context_set_context_device_id(device_id) + @property def feature_flags(self): """ diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index e9d7054c..74360add 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2294,6 +2294,35 @@ def test_device_id_from_context_is_used_in_flags_request(self, patch_flags): flag_keys_to_evaluate=["random_key"], ) + @mock.patch("posthog.client.flags") + def test_client_set_context_device_id_is_used_in_flags_request(self, patch_flags): + patch_flags.return_value = { + "featureFlags": { + "beta-feature": "random-variant", + } + } + client = Client( + FAKE_TEST_API_KEY, + on_error=self.set_fail, + ) + + with client.new_context(): + client.set_context_device_id("client-context-device-id") + client.get_feature_flag("random_key", "some_id") + + patch_flags.assert_called_with( + "random_key", + "https://us.i.posthog.com", + timeout=3, + distinct_id="some_id", + groups={}, + person_properties={"distinct_id": "some_id"}, + group_properties={}, + geoip_disable=True, + device_id="client-context-device-id", + flag_keys_to_evaluate=["random_key"], + ) + @parameterized.expand( [ # name, sys_platform, version_info, expected_runtime, expected_version, expected_os, expected_os_version, expected_os_distro, platform_method, platform_return @@ -2467,6 +2496,58 @@ def test_set_context_session_with_capture(self): msg["properties"]["$session_id"], "context-session-123" ) + def test_client_context_helpers_apply_to_capture(self): + with mock.patch("posthog.client.batch_post") as mock_post: + client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) + + with client.new_context(fresh=True): + client.tag("client_tag", "tag-value") + client.identify_context("context-user") + client.set_context_session("context-session-123") + + self.assertEqual(client.get_tags(), {"client_tag": "tag-value"}) + + msg_uuid = client.capture( + "test_event", + properties={"custom_prop": "value"}, + ) + + self.assertIsNotNone(msg_uuid) + mock_post.assert_called_once() + batch_data = mock_post.call_args[1]["batch"] + msg = batch_data[0] + + self.assertEqual(msg["distinct_id"], "context-user") + self.assertEqual(msg["properties"]["client_tag"], "tag-value") + self.assertEqual(msg["properties"]["custom_prop"], "value") + self.assertEqual(msg["properties"]["$session_id"], "context-session-123") + self.assertCountEqual(msg["properties"]["$context_tags"], ["client_tag"]) + self.assertEqual(client.get_tags(), {}) + + def test_client_scoped_context_helpers_apply_to_capture(self): + with mock.patch("posthog.client.batch_post") as mock_post: + client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) + + @client.scoped(fresh=True) + def capture_in_client_scope(): + client.tag("scoped_tag", "scoped-value") + client.identify_context("scoped-user") + client.set_context_session("scoped-session-123") + return client.capture("scoped_event") + + msg_uuid = capture_in_client_scope() + + self.assertIsNotNone(msg_uuid) + mock_post.assert_called_once() + batch_data = mock_post.call_args[1]["batch"] + msg = batch_data[0] + + self.assertEqual(msg["distinct_id"], "scoped-user") + self.assertEqual(msg["properties"]["scoped_tag"], "scoped-value") + self.assertEqual(msg["properties"]["$session_id"], "scoped-session-123") + self.assertCountEqual(msg["properties"]["$context_tags"], ["scoped_tag"]) + self.assertEqual(client.get_tags(), {}) + def test_set_context_session_with_page_explicit_properties(self): with mock.patch("posthog.client.batch_post") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 2ba8a253..1f6c9284 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -1059,13 +1059,19 @@ method posthog.client.Client.get_feature_payloads(distinct_id, groups=None, pers method posthog.client.Client.get_feature_variants(distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None, flag_keys_to_evaluate: Optional[list[str]] = None, device_id: Optional[str] = None) -> dict[str, Union[bool, str]] method posthog.client.Client.get_flags_decision(distinct_id: Optional[ID_TYPES] = None, groups: Optional[dict] = None, person_properties=None, group_properties=None, disable_geoip=None, flag_keys_to_evaluate: Optional[list[str]] = None, device_id: Optional[str] = None) -> FlagsResponse method posthog.client.Client.get_remote_config_payload(key: str) +method posthog.client.Client.get_tags() -> Dict[str, Any] method posthog.client.Client.group_identify(group_type: str, group_key: str, properties: Optional[Dict[str, Any]] = None, timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, distinct_id: Optional[ID_TYPES] = None) -> Optional[str] +method posthog.client.Client.identify_context(distinct_id: str) -> None method posthog.client.Client.join() -> None method posthog.client.Client.load_feature_flags() method posthog.client.Client.new_context(fresh=False, capture_exceptions=True) +method posthog.client.Client.scoped(fresh=False, capture_exceptions=True) method posthog.client.Client.set(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str] +method posthog.client.Client.set_context_device_id(device_id: str) -> None +method posthog.client.Client.set_context_session(session_id: str) -> None method posthog.client.Client.set_once(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str] method posthog.client.Client.shutdown() -> None +method posthog.client.Client.tag(name: str, value: Any) -> None method posthog.consumer.Consumer.next() method posthog.consumer.Consumer.pause() method posthog.consumer.Consumer.request(batch) From e51a1d9ee9a345244f912117358d7edad815a566 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Fri, 19 Jun 2026 14:57:34 +0200 Subject: [PATCH 2/4] address pr review feedback --- posthog/client.py | 32 ++++---------------- posthog/contexts.py | 15 ++++++++-- posthog/test/test_client.py | 48 +++++++++++++++++++++--------- references/public_api_snapshot.txt | 2 +- 4 files changed, 52 insertions(+), 45 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index 6ceb7850..8f9d3a32 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -8,7 +8,7 @@ import warnings import weakref from datetime import datetime, timedelta, timezone -from typing import Any, Callable, Dict, List, Optional, TypeVar, Union, cast +from typing import Any, Dict, List, Optional, Union from uuid import uuid4 from typing_extensions import Unpack @@ -27,6 +27,7 @@ get_tags as _context_get_tags, identify_context as _context_identify_context, new_context, + scoped as _context_scoped, set_context_device_id as _context_set_context_device_id, set_context_session as _context_set_context_session, tag as _context_tag, @@ -104,8 +105,6 @@ MAX_DICT_SIZE = 50_000 -_F = TypeVar("_F", bound=Callable[..., Any]) - def get_identity_state(passed) -> tuple[str, bool]: """Returns the distinct id to use, and whether this is a personless event or not""" @@ -497,30 +496,9 @@ def scoped(self, fresh=False, capture_exceptions=True): Contexts """ - def decorator(func: _F) -> _F: - from functools import wraps - - if inspect.iscoroutinefunction(func): - - @wraps(func) - async def async_wrapper(*args, **kwargs): - with self.new_context( - fresh=fresh, capture_exceptions=capture_exceptions - ): - return await func(*args, **kwargs) - - return cast(_F, async_wrapper) - - @wraps(func) - def wrapper(*args, **kwargs): - with self.new_context( - fresh=fresh, capture_exceptions=capture_exceptions - ): - return func(*args, **kwargs) - - return cast(_F, wrapper) - - return decorator + return _context_scoped( + fresh=fresh, capture_exceptions=capture_exceptions, client=self + ) def tag(self, name: str, value: Any) -> None: """ diff --git a/posthog/contexts.py b/posthog/contexts.py index 56bed4a9..f14624a9 100644 --- a/posthog/contexts.py +++ b/posthog/contexts.py @@ -370,7 +370,11 @@ def get_code_variables_ignore_patterns_context() -> Optional[list]: F = TypeVar("F", bound=Callable[..., Any]) -def scoped(fresh: bool = False, capture_exceptions: bool = True): +def scoped( + fresh: bool = False, + capture_exceptions: bool = True, + client: Optional["Client"] = None, +): """ Decorator that creates a new context for the function. Simply wraps the function in a with posthog.new_context(): block. @@ -378,6 +382,7 @@ def scoped(fresh: bool = False, capture_exceptions: bool = True): Args: fresh: Whether to start with a fresh context (default: False) capture_exceptions: Whether to capture and track exceptions with posthog error tracking (default: True) + client: Optional client instance to use for capturing exceptions (default: None) Example: @posthog.scoped() @@ -410,14 +415,18 @@ def decorator(func: F) -> F: @wraps(func) async def async_wrapper(*args, **kwargs): - with new_context(fresh=fresh, capture_exceptions=capture_exceptions): + with new_context( + fresh=fresh, capture_exceptions=capture_exceptions, client=client + ): return await func(*args, **kwargs) return cast(F, async_wrapper) @wraps(func) def wrapper(*args, **kwargs): - with new_context(fresh=fresh, capture_exceptions=capture_exceptions): + with new_context( + fresh=fresh, capture_exceptions=capture_exceptions, client=client + ): return func(*args, **kwargs) return cast(F, wrapper) diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 74360add..e01c3300 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -1,3 +1,4 @@ +import asyncio import time import unittest from datetime import datetime @@ -2496,22 +2497,34 @@ def test_set_context_session_with_capture(self): msg["properties"]["$session_id"], "context-session-123" ) - def test_client_context_helpers_apply_to_capture(self): + @parameterized.expand([("new_context",), ("scoped",)]) + def test_client_context_helpers_apply_to_capture(self, context_helper): with mock.patch("posthog.client.batch_post") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) - with client.new_context(fresh=True): + def capture_in_context(): client.tag("client_tag", "tag-value") client.identify_context("context-user") client.set_context_session("context-session-123") self.assertEqual(client.get_tags(), {"client_tag": "tag-value"}) - msg_uuid = client.capture( + return client.capture( "test_event", properties={"custom_prop": "value"}, ) + if context_helper == "new_context": + with client.new_context(fresh=True): + msg_uuid = capture_in_context() + else: + + @client.scoped(fresh=True) + def scoped_capture(): + return capture_in_context() + + msg_uuid = scoped_capture() + self.assertIsNotNone(msg_uuid) mock_post.assert_called_once() batch_data = mock_post.call_args[1]["batch"] @@ -2524,28 +2537,35 @@ def test_client_context_helpers_apply_to_capture(self): self.assertCountEqual(msg["properties"]["$context_tags"], ["client_tag"]) self.assertEqual(client.get_tags(), {}) - def test_client_scoped_context_helpers_apply_to_capture(self): + def test_client_scoped_context_helpers_apply_to_capture_async(self): with mock.patch("posthog.client.batch_post") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) @client.scoped(fresh=True) - def capture_in_client_scope(): - client.tag("scoped_tag", "scoped-value") - client.identify_context("scoped-user") - client.set_context_session("scoped-session-123") - return client.capture("scoped_event") + async def scoped_capture(): + client.tag("async_scoped_tag", "async-scoped-value") + client.identify_context("async-scoped-user") + client.set_context_session("async-scoped-session-123") + await asyncio.sleep(0) + return client.capture("async_scoped_event") - msg_uuid = capture_in_client_scope() + msg_uuid = asyncio.run(scoped_capture()) self.assertIsNotNone(msg_uuid) mock_post.assert_called_once() batch_data = mock_post.call_args[1]["batch"] msg = batch_data[0] - self.assertEqual(msg["distinct_id"], "scoped-user") - self.assertEqual(msg["properties"]["scoped_tag"], "scoped-value") - self.assertEqual(msg["properties"]["$session_id"], "scoped-session-123") - self.assertCountEqual(msg["properties"]["$context_tags"], ["scoped_tag"]) + self.assertEqual(msg["distinct_id"], "async-scoped-user") + self.assertEqual( + msg["properties"]["async_scoped_tag"], "async-scoped-value" + ) + self.assertEqual( + msg["properties"]["$session_id"], "async-scoped-session-123" + ) + self.assertCountEqual( + msg["properties"]["$context_tags"], ["async_scoped_tag"] + ) self.assertEqual(client.get_tags(), {}) def test_set_context_session_with_page_explicit_properties(self): diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 1f6c9284..baafa162 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -871,7 +871,7 @@ function posthog.contexts.get_context_session_id() -> Optional[str] function posthog.contexts.get_tags() -> Dict[str, Any] function posthog.contexts.identify_context(distinct_id: str) -> None function posthog.contexts.new_context(fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) -function posthog.contexts.scoped(fresh: bool = False, capture_exceptions: bool = True) +function posthog.contexts.scoped(fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) function posthog.contexts.set_capture_exception_code_variables_context(enabled: bool) -> None function posthog.contexts.set_code_variables_ignore_patterns_context(ignore_patterns: list) -> None function posthog.contexts.set_code_variables_mask_patterns_context(mask_patterns: list) -> None From 3139c870f3dbdf81fc49ebb7ed7aa26e8046d9b0 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Fri, 19 Jun 2026 15:06:13 +0200 Subject: [PATCH 3/4] keep scoped public api stable --- posthog/client.py | 2 +- posthog/contexts.py | 62 +++++++++++++++--------------- references/public_api_snapshot.txt | 2 +- 3 files changed, 32 insertions(+), 34 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index 8f9d3a32..ac02eab2 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -26,8 +26,8 @@ get_context_session_id, get_tags as _context_get_tags, identify_context as _context_identify_context, + _scoped as _context_scoped, new_context, - scoped as _context_scoped, set_context_device_id as _context_set_context_device_id, set_context_session as _context_set_context_session, tag as _context_tag, diff --git a/posthog/contexts.py b/posthog/contexts.py index f14624a9..cd64941c 100644 --- a/posthog/contexts.py +++ b/posthog/contexts.py @@ -370,11 +370,35 @@ def get_code_variables_ignore_patterns_context() -> Optional[list]: F = TypeVar("F", bound=Callable[..., Any]) -def scoped( - fresh: bool = False, - capture_exceptions: bool = True, - client: Optional["Client"] = None, -): +def _scoped(fresh: bool = False, capture_exceptions: bool = True, client=None): + def decorator(func: F) -> F: + from functools import wraps + from inspect import iscoroutinefunction + + if iscoroutinefunction(func): + + @wraps(func) + async def async_wrapper(*args, **kwargs): + with new_context( + fresh=fresh, capture_exceptions=capture_exceptions, client=client + ): + return await func(*args, **kwargs) + + return cast(F, async_wrapper) + + @wraps(func) + def wrapper(*args, **kwargs): + with new_context( + fresh=fresh, capture_exceptions=capture_exceptions, client=client + ): + return func(*args, **kwargs) + + return cast(F, wrapper) + + return decorator + + +def scoped(fresh: bool = False, capture_exceptions: bool = True): """ Decorator that creates a new context for the function. Simply wraps the function in a with posthog.new_context(): block. @@ -382,7 +406,6 @@ def scoped( Args: fresh: Whether to start with a fresh context (default: False) capture_exceptions: Whether to capture and track exceptions with posthog error tracking (default: True) - client: Optional client instance to use for capturing exceptions (default: None) Example: @posthog.scoped() @@ -406,29 +429,4 @@ async def middleware(request, call_next): Category: Contexts """ - - def decorator(func: F) -> F: - from functools import wraps - from inspect import iscoroutinefunction - - if iscoroutinefunction(func): - - @wraps(func) - async def async_wrapper(*args, **kwargs): - with new_context( - fresh=fresh, capture_exceptions=capture_exceptions, client=client - ): - return await func(*args, **kwargs) - - return cast(F, async_wrapper) - - @wraps(func) - def wrapper(*args, **kwargs): - with new_context( - fresh=fresh, capture_exceptions=capture_exceptions, client=client - ): - return func(*args, **kwargs) - - return cast(F, wrapper) - - return decorator + return _scoped(fresh=fresh, capture_exceptions=capture_exceptions) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index baafa162..1f6c9284 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -871,7 +871,7 @@ function posthog.contexts.get_context_session_id() -> Optional[str] function posthog.contexts.get_tags() -> Dict[str, Any] function posthog.contexts.identify_context(distinct_id: str) -> None function posthog.contexts.new_context(fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) -function posthog.contexts.scoped(fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) +function posthog.contexts.scoped(fresh: bool = False, capture_exceptions: bool = True) function posthog.contexts.set_capture_exception_code_variables_context(enabled: bool) -> None function posthog.contexts.set_code_variables_ignore_patterns_context(ignore_patterns: list) -> None function posthog.contexts.set_code_variables_mask_patterns_context(mask_patterns: list) -> None From 75778523a7f9a7cde034ed1c72877609c4e573d0 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Fri, 19 Jun 2026 15:09:01 +0200 Subject: [PATCH 4/4] mark client context helpers as minor --- .changeset/fuzzy-pandas-scope.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.changeset/fuzzy-pandas-scope.md b/.changeset/fuzzy-pandas-scope.md index 25f9053d..7e1e359c 100644 --- a/.changeset/fuzzy-pandas-scope.md +++ b/.changeset/fuzzy-pandas-scope.md @@ -1,5 +1,5 @@ --- -'pypi/posthog': patch +'pypi/posthog': minor --- Add context helper methods to custom PostHog client instances.