From 6d7f522e1cc5181e75209e4c00109426baa335fc Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 3 Dec 2023 15:58:39 -1000
Subject: [PATCH] feat: speed up sending messages with call on the MessageBus
 (#271)

---
 src/dbus_fast/aio/message_bus.py  |  2 +-
 src/dbus_fast/glib/message_bus.py |  1 +
 src/dbus_fast/message_bus.pxd     |  5 +++++
 src/dbus_fast/message_bus.py      | 36 +++++++++++++++----------------
 4 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/src/dbus_fast/aio/message_bus.py b/src/dbus_fast/aio/message_bus.py
index 46b2e9d..620152e 100644
--- a/src/dbus_fast/aio/message_bus.py
+++ b/src/dbus_fast/aio/message_bus.py
@@ -394,7 +394,7 @@ class MessageBus(BaseMessageBus):
                 else:
                     _future_set_result(future, reply)
 
-        self._call(msg, reply_handler, check_callback=False)
+        self._call(msg, reply_handler)
 
         await future
 
diff --git a/src/dbus_fast/glib/message_bus.py b/src/dbus_fast/glib/message_bus.py
index ece6f97..42de20d 100644
--- a/src/dbus_fast/glib/message_bus.py
+++ b/src/dbus_fast/glib/message_bus.py
@@ -288,6 +288,7 @@ class MessageBus(BaseMessageBus):
             this message. May return an :class:`Exception` on connection errors.
         :type reply_notify: Callable
         """
+        BaseMessageBus._check_callback_type(reply_notify)
         self._call(msg, reply_notify)
 
     def call_sync(self, msg: Message) -> Optional[Message]:
diff --git a/src/dbus_fast/message_bus.pxd b/src/dbus_fast/message_bus.pxd
index 2114e6f..58e883b 100644
--- a/src/dbus_fast/message_bus.pxd
+++ b/src/dbus_fast/message_bus.pxd
@@ -52,3 +52,8 @@ cdef class BaseMessageBus:
     cdef _find_message_handler(self, Message msg)
 
     cdef _setup_socket(self)
+
+    @cython.locals(no_reply_expected=bint)
+    cpdef _call(self, Message msg, object callback)
+
+    cpdef next_serial(self)
diff --git a/src/dbus_fast/message_bus.py b/src/dbus_fast/message_bus.py
index 8dd00c0..6bb71ba 100644
--- a/src/dbus_fast/message_bus.py
+++ b/src/dbus_fast/message_bus.py
@@ -3,6 +3,7 @@ import logging
 import socket
 import traceback
 import xml.etree.ElementTree as ET
+from functools import partial
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 from . import introspection as intr
@@ -713,34 +714,35 @@ class BaseMessageBus:
         if err:
             raise err
 
+    def _reply_notify(
+        self,
+        msg: Message,
+        callback: Optional[Callable[[Optional[Message], Optional[Exception]], None]],
+        reply: Optional[Message],
+        err: Optional[Exception],
+    ) -> None:
+        """Callback on reply."""
+        if reply and msg.destination and reply.sender:
+            self._name_owners[msg.destination] = reply.sender
+        callback(reply, err)
+
     def _call(
         self,
         msg: Message,
         callback: Optional[Callable[[Optional[Message], Optional[Exception]], None]],
-        check_callback: bool = True,
     ) -> None:
-        if check_callback:
-            BaseMessageBus._check_callback_type(callback)
-
         if not msg.serial:
             msg.serial = self.next_serial()
 
-        no_reply_expected = _expects_reply(msg) is False
+        no_reply_expected = not _expects_reply(msg)
         # Make sure the return reply handler is installed
         # before sending the message to avoid a race condition
         # where the reply is lost in case the backend can
         # send it right away.
         if not no_reply_expected:
-
-            def _reply_notify(
-                reply: Optional[Message], err: Optional[Exception]
-            ) -> None:
-                """Callback on reply."""
-                if reply and msg.destination and reply.sender:
-                    self._name_owners[msg.destination] = reply.sender
-                callback(reply, err)
-
-            self._method_return_handlers[msg.serial] = _reply_notify
+            self._method_return_handlers[msg.serial] = partial(
+                self._reply_notify, msg, callback
+            )
 
         self.send(msg)
 
@@ -986,7 +988,6 @@ class BaseMessageBus:
                 member="GetMachineId",
             ),
             reply_handler,
-            check_callback=False,
         )
 
     def _default_get_managed_objects_handler(
@@ -1213,7 +1214,6 @@ class BaseMessageBus:
                 body=[self._name_owner_match_rule],
             ),
             add_match_notify,
-            check_callback=False,
         )
 
     def _add_match_rule(self, match_rule):
@@ -1247,7 +1247,6 @@ class BaseMessageBus:
                 body=[match_rule],
             ),
             add_match_notify,
-            check_callback=False,
         )
 
     def _remove_match_rule(self, match_rule):
@@ -1286,5 +1285,4 @@ class BaseMessageBus:
                 body=[match_rule],
             ),
             remove_match_notify,
-            check_callback=False,
         )