feat: session and identity integrate with context now

CHANGELOG.md

@@ -1,4 +1,9 @@
-## 5.0.0
+# 5.1.0
+
+- feat: session and distinct ID's can now be associated with contexts, and are used as such
+- feat: django http request middleware
+
+## 5.0.0 - 2025-06-16
 
 - fix: removed deprecated sentry integration
 

posthog/__init__.py

@@ -4,7 +4,15 @@
 
 from posthog.client import Client
 from posthog.exception_capture import Integrations  # noqa: F401
-from posthog.scopes import clear_tags, get_tags, new_context, scoped, tag
+from posthog.scopes import (
+    clear_tags,
+    get_tags,
+    new_context,
+    scoped,
+    tag,
+    set_context_session,
+    identify_context,
+)
 from posthog.types import FeatureFlag, FlagsAndPayloads
 from posthog.version import VERSION
 
@@ -16,6 +24,9 @@
 get_tags = get_tags
 clear_tags = clear_tags
 scoped = scoped
+identify_context = identify_context
+set_context_session = set_context_session
+
 
 """Settings."""
 api_key = None  # type: Optional[str]

posthog/client.py

@@ -33,7 +33,11 @@
     get,
     remote_config,
 )
-from posthog.scopes import get_tags
+from posthog.scopes import (
+    _get_current_context,
+    get_context_distinct_id,
+    get_context_session_id,
+)
 from posthog.types import (
     FeatureFlag,
     FeatureFlagResult,
@@ -285,10 +289,17 @@ def identify(
                 stacklevel=2,
             )
 
+        if distinct_id is None:
+            distinct_id = get_context_distinct_id()
+
         properties = properties or {}
+
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
+        if "$session_id" not in properties and get_context_session_id():
+            properties["$session_id"] = get_context_session_id()
+
         msg = {
             "timestamp": timestamp,
             "distinct_id": distinct_id,
@@ -358,6 +369,9 @@ def get_flags_decision(
         """
         Get feature flags decision, using either flags() or decide() API based on rollout.
         """
+
+        if distinct_id is None:
+            distinct_id = get_context_distinct_id()
         require("distinct_id", distinct_id, ID_TYPES)
 
         if disable_geoip is None:
@@ -406,14 +420,19 @@ def capture(
 
         properties = {**(properties or {}), **system_context()}
 
+        if "$session_id" not in properties and get_context_session_id():
+            properties["$session_id"] = get_context_session_id()
+
+        if distinct_id is None:
+            distinct_id = get_context_distinct_id()
+
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
         require("event", event, string_types)
 
-        # Grab current context tags, if any exist
-        context_tags = get_tags()
-        if context_tags:
-            properties.update(context_tags)
+        current_context = _get_current_context()
+        if current_context:
+            properties.update(current_context.collect_tags())
 
         msg = {
             "properties": properties,
@@ -480,6 +499,9 @@ def set(
                 stacklevel=2,
             )
 
+        if distinct_id is None:
+            distinct_id = get_context_distinct_id()
+
         properties = properties or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
@@ -510,6 +532,9 @@ def set_once(
                 stacklevel=2,
             )
 
+        if distinct_id is None:
+            distinct_id = get_context_distinct_id()
+
         properties = properties or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
@@ -581,6 +606,9 @@ def alias(
                 stacklevel=2,
             )
 
+        if distinct_id is None:
+            distinct_id = get_context_distinct_id()
+
         require("previous_id", previous_id, ID_TYPES)
         require("distinct_id", distinct_id, ID_TYPES)
 
@@ -613,6 +641,9 @@ def page(
                 stacklevel=2,
             )
 
+        if distinct_id is None:
+            distinct_id = get_context_distinct_id()
+
         properties = properties or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
@@ -648,6 +679,9 @@ def capture_exception(
                 stacklevel=2,
             )
 
+        if distinct_id is None:
+            distinct_id = get_context_distinct_id()
+
         # this function shouldn't ever throw an error, so it logs exceptions instead of raising them.
         # this is important to ensure we don't unexpectedly re-raise exceptions in the user's code.
         try:

posthog/scopes.py

@@ -1,14 +1,66 @@
 import contextvars
 from contextlib import contextmanager
-from typing import Any, Callable, Dict, TypeVar, cast
-
-_context_stack: contextvars.ContextVar[list] = contextvars.ContextVar(
-    "posthog_context_stack", default=[{}]
+from typing import Self, Optional, Any, Callable, Dict, TypeVar, cast
+
+
+class ContextScope:
+    def __init__(
+        self,
+        parent: Optional[Self],
+        fresh: bool = False,
+        capture_exceptions: bool = True,
+    ):
+        self.parent = parent
+        self.fresh = fresh
+        self.capture_exceptions = capture_exceptions
+        self.session_id: Optional[str] = None
+        self.distinct_id: Optional[str] = None
+        self.tags: Dict[str, Any] = {}
+
+    def set_session_id(self, session_id: str):
+        self.session_id = session_id
+
+    def set_distinct_id(self, distinct_id: str):
+        self.distinct_id = distinct_id
+
+    def add_tag(self, key: str, value: Any):
+        self.tags[key] = value
+
+    def get_parent(self) -> Optional[Self]:
+        return self.parent
+
+    def get_session_id(self) -> Optional[str]:
+        if self.session_id is not None:
+            return self.session_id
+        if self.parent is not None and not self.fresh:
+            return self.parent.get_session_id()
+        return None
+
+    def get_distinct_id(self) -> Optional[str]:
+        if self.distinct_id is not None:
+            return self.distinct_id
+        if self.parent is not None and not self.fresh:
+            return self.parent.get_distinct_id()
+        return None
+
+    def collect_tags(self) -> Dict[str, Any]:
+        tags = self.tags.copy()
+        if self.parent and not self.fresh:
+            # We want child tags to take precedence over parent tags,
+            # so we can't use a simple update here, instead collecting
+            # the parent tags and then updating with the child tags.
+            new_tags = self.parent.collect_tags()
+            tags.update(new_tags)
+        return tags
+
+
+_context_stack: contextvars.ContextVar[Optional[ContextScope]] = contextvars.ContextVar(
+    "posthog_context_stack", default=None
 )
 
 
-def _get_current_context() -> Dict[str, Any]:
-    return _context_stack.get()[-1]
+def _get_current_context() -> Optional[ContextScope]:
+    return _context_stack.get()
 
 
 @contextmanager
@@ -44,19 +96,18 @@ def new_context(fresh=False, capture_exceptions=True):
     """
     from posthog import capture_exception
 
-    current_tags = _get_current_context().copy()
-    current_stack = _context_stack.get()
-    new_stack = current_stack + [{}] if fresh else current_stack + [current_tags]
-    token = _context_stack.set(new_stack)
+    current_context = _get_current_context()
+    new_context = ContextScope(current_context, fresh, capture_exceptions)
+    _context_stack.set(new_context)
 
     try:
         yield
     except Exception as e:
-        if capture_exceptions:
+        if new_context.capture_exceptions:
             capture_exception(e)
         raise
     finally:
-        _context_stack.reset(token)
+        _context_stack.set(new_context.get_parent())
 
 
 def tag(key: str, value: Any) -> None:
@@ -70,9 +121,13 @@ def tag(key: str, value: Any) -> None:
     Example:
         posthog.tag("user_id", "123")
     """
-    _get_current_context()[key] = value
+    current_context = _get_current_context()
+    if current_context:
+        current_context.add_tag(key, value)
 
 
+# NOTE: we should probably also remove this - there's no reason for the user to ever
+# need to manually interact with the current tag set
 def get_tags() -> Dict[str, Any]:
     """
     Get all tags from the current context. Note, modifying
@@ -81,12 +136,73 @@ def get_tags() -> Dict[str, Any]:
     Returns:
         Dict of all tags in the current context
     """
-    return _get_current_context().copy()
+    current_context = _get_current_context()
+    if current_context:
+        return current_context.collect_tags()
+    return {}
 
 
+# NOTE: We should probably remove this function - the way to clear scope context
+# is by entering a new, fresh context, rather than by clearing the tags or other
+# scope data directly.
 def clear_tags() -> None:
-    """Clear all tags in the current context."""
-    _get_current_context().clear()
+    """Clear all tags in the current context. Does not clear parent tags"""
+    current_context = _get_current_context()
+    if current_context:
+        current_context.tags.clear()
+
+
+def identify_context(distinct_id: str) -> None:
+    """
+    Identify the current context with a distinct ID, associating all events captured in this or
+    child contexts with the given distinct ID (unless identify_scope is called again). This is overridden by
+    distinct id's passed directly to posthog.capture and related methods (identify, set etc).
+
+    Args:
+        distinct_id: The distinct ID to associate with the current context and its children.
+    """
+    current_context = _get_current_context()
+    if current_context:
+        current_context.set_distinct_id(distinct_id)
+
+
+def set_context_session(session_id: str) -> None:
+    """
+    Set the session ID for the current context, associating all events captured in this or
+    child contexts with the given session ID (unless set_scope_session is called again).
+
+    Args:
+        session_id: The session ID to associate with the current context and its children. See https://posthog.com/docs/data/sessions
+    """
+    current_context = _get_current_context()
+    if current_context:
+        current_context.set_session_id(session_id)
+
+
+def get_context_session_id() -> Optional[str]:
+    """
+    Get the session ID for the current context.
+
+    Returns:
+        The session ID if set, None otherwise
+    """
+    current_context = _get_current_context()
+    if current_context:
+        return current_context.get_session_id()
+    return None
+
+
+def get_context_distinct_id() -> Optional[str]:
+    """
+    Get the distinct ID for the current context.
+
+    Returns:
+        The distinct ID if set, None otherwise
+    """
+    current_context = _get_current_context()
+    if current_context:
+        return current_context.get_distinct_id()
+    return None
 
 
 F = TypeVar("F", bound=Callable[..., Any])