DRAFT: Added get_current_message call

igo95862 · igo95862 · commit dffbfb005259 · 2022-07-21T20:58:48.000+03:00
diff --git a/src/sdbus/__init__.py b/src/sdbus/__init__.py
@@ -64,6 +64,7 @@
 from .dbus_proxy_async_method import (
     dbus_method_async,
     dbus_method_async_override,
+    get_current_message,
 )
 from .dbus_proxy_async_property import (
     dbus_property_async,
@@ -134,6 +135,7 @@
 
     'dbus_property_async',
     'dbus_property_async_override',
+    'get_current_message',
 
     'dbus_signal_async',
 
diff --git a/src/sdbus/dbus_proxy_async_method.py b/src/sdbus/dbus_proxy_async_method.py
@@ -19,6 +19,7 @@
 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
 from __future__ import annotations
 
+from contextvars import ContextVar, copy_context
 from inspect import iscoroutinefunction
 from types import FunctionType
 from typing import (
@@ -42,6 +43,13 @@
 from .dbus_exceptions import DbusFailedError
 from .sd_bus_internals import DbusNoReplyFlag, SdBusMessage
 
+CURRENT_MESSAGE: ContextVar[SdBusMessage] = ContextVar('CURRENT_MESSAGE')
+
+
+def get_current_message() -> SdBusMessage:
+    return CURRENT_MESSAGE.get()
+
+
 T_input = TypeVar('T_input')
 T = TypeVar('T')
 
@@ -118,24 +126,38 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
             return self.dbus_method.original_method(
                 interface, *args, **kwargs)
 
+    async def _call_method_from_dbus(
+            self,
+            request_message: SdBusMessage,
+            interface: DbusInterfaceBaseAsync) -> Any:
+        request_data = request_message.get_contents()
+
+        local_method = self.dbus_method.original_method.__get__(
+            interface, None)
+
+        CURRENT_MESSAGE.set(request_message)
+
+        if isinstance(request_data, tuple):
+            return await local_method(*request_data)
+        elif request_data is None:
+            return await local_method()
+        else:
+            return await local_method(request_data)
+
     async def _call_from_dbus(
             self,
             request_message: SdBusMessage) -> None:
         interface = self.interface_ref()
         assert interface is not None
 
-        request_data = request_message.get_contents()
-
-        local_method = self.dbus_method.original_method.__get__(
-            interface, None)
+        call_context = copy_context()
 
         try:
-            if isinstance(request_data, tuple):
-                reply_data = await local_method(*request_data)
-            elif request_data is None:
-                reply_data = await local_method()
-            else:
-                reply_data = await local_method(request_data)
+            reply_data = await call_context.run(
+                self._call_method_from_dbus,
+                request_message,
+                interface,
+            )
         except DbusFailedError as e:
             if not request_message.expect_reply:
                 return
diff --git a/test/test_sd_bus_async.py b/test/test_sd_bus_async.py
@@ -48,6 +48,7 @@
     dbus_property_async,
     dbus_property_async_override,
     dbus_signal_async,
+    get_current_message,
 )
 
 
@@ -98,6 +99,11 @@ async def upper(self, string: str) -> str:
         """Uppercase the input"""
         return string.upper()
 
+    @dbus_method_async(result_signature='s')
+    async def get_sender(self) -> str:
+        message = get_current_message()
+        return message.sender or ''
+
     @dbus_method_async(result_signature='x')
     async def test_int(self) -> int:
         return 1
@@ -281,6 +287,8 @@ async def test_method(self) -> None:
                 test_object_connection.test_struct_return_workaround(), 0.5),
         )
 
+        self.assertTrue(await test_object_connection.get_sender())
+
     async def test_subclass(self) -> None:
         test_object, test_object_connection = initialize_object()
 
