Add sdbus.default_bus.set_context_default_bus

igo95862 · igo95862 · commit 8564580d0df3 · 2025-04-06T18:56:37.000+01:00
Sets the context-local default bus. The context bus only used
when explicitly set to avoid bus initialization spam.

Only having thread-local buses seemed very limiting.
diff --git a/docs/general.rst b/docs/general.rst
@@ -130,7 +130,7 @@ See API documentation for a particular decorator.
 
 
 Default bus
-++++++++++++++++++++++++++
++++++++++++
 
 Most object methods that take a bus as a parameter
 will use a thread-local default bus connection if a bus object
@@ -139,17 +139,23 @@ is not explicitly passed.
 Session bus is default bus when running as a user and
 system bus otherwise.
 
-:py:func:`request_default_bus_name_async` can be used to acquire
-a service name on default bus.
+The :py:func:`request_default_bus_name_async <sdbus.default_bus.request_default_bus_name_async>`
+and :py:func:`request_default_bus_name <sdbus.default_bus.request_default_bus_name>`
+can be used to acquire a service name on the default bus.
 
 Use :py:func:`sd_bus_open_user` and :py:func:`sd_bus_open_system` to
 acquire a specific bus connection.
 
-Set the default connection to a new default with :py:func:`set_default_bus`.
-This should be done before any object that take bus as an init argument are created.
+The :py:func:`set_default_bus <sdbus.default_bus.set_default_bus>` can be used to set the new
+thread-local bus. This should be done before any objects that take bus as
+an init argument are created. If no bus has been set the new bus will
+be initialized and set as thread-local default.
 
-In the future there will be a better way to create and acquire
-new bus connections.
+The bus can also be set as default for the current context using
+:py:func:`set_context_default_bus <sdbus.default_bus.set_context_default_bus>`.
+The context refers to the standard library's ``contextvars`` module context variables
+frequently used in asyncio frameworks. Context-local default bus has higher priority over
+thread-local default bus.
 
 Glossary
 +++++++++++++++++++++
diff --git a/src/sdbus/__init__.py b/src/sdbus/__init__.py
@@ -74,6 +74,7 @@
     get_default_bus,
     request_default_bus_name,
     request_default_bus_name_async,
+    set_context_default_bus,
     set_default_bus,
 )
 from .sd_bus_internals import (
@@ -145,6 +146,7 @@
     "get_default_bus",
     "request_default_bus_name",
     "request_default_bus_name_async",
+    "set_context_default_bus",
     "set_default_bus",
 
     'DbusDeprecatedFlag',
diff --git a/src/sdbus/default_bus.py b/src/sdbus/default_bus.py
@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 import threading
+from contextvars import ContextVar, Token
 from logging import getLogger
 from typing import TYPE_CHECKING
 
@@ -43,6 +44,7 @@ class DefaultBusTLStorage(threading.local):
 
 
 bus_tls = DefaultBusTLStorage()
+bus_contextvar: ContextVar[SdBus] = ContextVar("DEFAULT_BUS")
 
 
 def _get_defaul_bus_tls() -> Optional[SdBus]:
@@ -54,10 +56,20 @@ def _set_default_bus_tls(new_bus: Optional[SdBus]) -> None:
 
 
 def get_default_bus() -> SdBus:
-    """Get default thread-local bus."""
-    current_bus = _get_defaul_bus_tls()
-    if current_bus is not None:
-        return current_bus
+    """Get default bus.
+
+    Returns context-local default bus if set or
+    thread-local otherwise.
+
+    If no default bus is set initializes a new bus using
+    :py:func:`sdbus.sd_bus_open` and sets it as a thread-local
+    default bus.
+    """
+    if (context_bus := bus_contextvar.get(None)) is not None:
+        return context_bus
+
+    if (tls_bus := _get_defaul_bus_tls()) is not None:
+        return tls_bus
     else:
         new_bus = sd_bus_open()
         logger.info(
@@ -69,7 +81,7 @@ def get_default_bus() -> SdBus:
 
 
 def set_default_bus(new_default: SdBus) -> None:
-    """Set default thread-local bus.
+    """Set thread-local default bus.
 
     Should be called before creating any objects that will use
     default bus.
@@ -80,6 +92,25 @@ def set_default_bus(new_default: SdBus) -> None:
     _set_default_bus_tls(new_default)
 
 
+def set_context_default_bus(new_default: SdBus) -> Token[SdBus]:
+    """Set context-local default bus.
+
+    Should be called before creating any objects that will use
+    default bus.
+
+    Default bus can be replaced but the change will only affect
+    newly created objects.
+
+    Context-local default bus has higher priority over thread-local one
+    but has to be explicitly set.
+
+    :returns:
+        Token that can be used to reset context bus back.
+        See ``contextvars`` documentation for details.
+    """
+    return bus_contextvar.set(new_default)
+
+
 def _prepare_request_name_flags(
         allow_replacement: bool,
         replace_existing: bool,
@@ -167,6 +198,7 @@ def request_default_bus_name(
 __all__ = (
     "get_default_bus",
     "set_default_bus",
+    "set_context_default_bus",
     "request_default_bus_name_async",
     "request_default_bus_name",
 )
