From ef7d9d440775cf0ddcb6b3bc6115b3884be35792 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Fri, 4 Nov 2022 11:37:12 +0100
Subject: [PATCH] feat: use cimports for message marshalling (#149)

---
 src/dbus_fast/_private/marshaller.pxd   |  6 ++++++
 src/dbus_fast/_private/marshaller.py    |  7 +++++++
 src/dbus_fast/_private/unmarshaller.pxd |  3 +++
 src/dbus_fast/message.pxd               | 17 ++++++++++++-----
 src/dbus_fast/message.py                | 15 +++++++++------
 5 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/src/dbus_fast/_private/marshaller.pxd b/src/dbus_fast/_private/marshaller.pxd
index 3a8f264..f94c95d 100644
--- a/src/dbus_fast/_private/marshaller.pxd
+++ b/src/dbus_fast/_private/marshaller.pxd
@@ -9,12 +9,16 @@ cdef bytes PACKED_UINT32_ZERO
 cdef bytes PACKED_BOOL_TRUE
 cdef bytes PACKED_BOOL_FALSE
 
+cdef get_signature_tree
+
 cdef class Marshaller:
 
     cdef object signature_tree
     cdef bytearray _buf
     cdef object body
 
+    cdef _buffer(self)
+
     cpdef align(self, unsigned int n)
 
     @cython.locals(
@@ -89,6 +93,8 @@ cdef class Marshaller:
 
     cpdef marshall(self)
 
+    cdef _marshall(self)
+
     @cython.locals(
         offset=cython.ulong,
         t=cython.str,
diff --git a/src/dbus_fast/_private/marshaller.py b/src/dbus_fast/_private/marshaller.py
index 42dfa89..c269d46 100644
--- a/src/dbus_fast/_private/marshaller.py
+++ b/src/dbus_fast/_private/marshaller.py
@@ -24,6 +24,9 @@ class Marshaller:
     def buffer(self) -> bytearray:
         return self._buf
 
+    def _buffer(self) -> bytearray:
+        return self._buf
+
     def align(self, n):
         return self._align(n)
 
@@ -168,6 +171,10 @@ class Marshaller:
             return writer(self, body, type_)
 
     def marshall(self) -> bytearray:
+        """Marshalls the body into a byte array"""
+        return self._marshall()
+
+    def _marshall(self) -> bytearray:
         """Marshalls the body into a byte array"""
         try:
             self._construct_buffer()
diff --git a/src/dbus_fast/_private/unmarshaller.pxd b/src/dbus_fast/_private/unmarshaller.pxd
index 667f8d7..18c9ff8 100644
--- a/src/dbus_fast/_private/unmarshaller.pxd
+++ b/src/dbus_fast/_private/unmarshaller.pxd
@@ -72,6 +72,9 @@ cdef unsigned int TOKEN_G_AS_INT
 
 cdef object MARSHALL_STREAM_END_ERROR
 
+cdef get_signature_tree
+
+
 cdef inline unsigned long _cast_uint32_native(const char * payload, unsigned int offset):
     cdef unsigned long *u32p = <unsigned long *> &payload[offset]
     return u32p[0]
diff --git a/src/dbus_fast/message.pxd b/src/dbus_fast/message.pxd
index 21e35c2..051007c 100644
--- a/src/dbus_fast/message.pxd
+++ b/src/dbus_fast/message.pxd
@@ -2,13 +2,14 @@
 
 import cython
 
+from ._private.marshaller cimport Marshaller
+from .signature cimport Variant
+
 
 cdef object ErrorType
 cdef object SignatureTree
 cdef object SignatureType
 cdef object MessageType
-cdef object Variant
-cdef object Marshaller
 
 
 cdef object HEADER_PATH
@@ -25,6 +26,8 @@ cdef object HEADER_UNIX_FDS
 cdef object LITTLE_ENDIAN
 cdef object PROTOCOL_VERSION
 
+cdef object MESSAGE_FLAG
+
 cdef get_signature_tree
 
 cdef class Message:
@@ -36,12 +39,16 @@ cdef class Message:
     cdef public object message_type
     cdef public object flags
     cdef public object error_name
-    cdef public unsigned int reply_serial
+    cdef public object reply_serial
     cdef public object sender
     cdef public object unix_fds
     cdef public object signature
     cdef public object signature_tree
     cdef public object body
-    cdef public unsigned int serial
+    cdef public object serial
 
-    cpdef _marshall(self, bint negotiate_unix_fd)
+    @cython.locals(
+        body_buffer=cython.bytearray,
+        header_buffer=cython.bytearray
+    )
+    cpdef _marshall(self, object negotiate_unix_fd)
diff --git a/src/dbus_fast/message.py b/src/dbus_fast/message.py
index 768973d..ccb2cba 100644
--- a/src/dbus_fast/message.py
+++ b/src/dbus_fast/message.py
@@ -28,6 +28,8 @@ HEADER_DESTINATION = HeaderField.DESTINATION.value
 HEADER_SIGNATURE = HeaderField.SIGNATURE.value
 HEADER_UNIX_FDS = HeaderField.UNIX_FDS.value
 
+MESSAGE_FLAG = MessageFlag
+
 
 class Message:
     """A class for sending and receiving messages through the
@@ -115,7 +117,7 @@ class Message:
         self.interface = interface
         self.member = member
         self.message_type = message_type
-        self.flags = flags if type(flags) is MessageFlag else MessageFlag(flags)
+        self.flags = flags if type(flags) is MESSAGE_FLAG else MESSAGE_FLAG(flags)
         self.error_name = (
             str(error_name.value) if type(error_name) is ErrorType else error_name
         )
@@ -258,7 +260,7 @@ class Message:
         """Marshall this message into a byte array."""
         # TODO maximum message size is 134217728 (128 MiB)
         body_block = Marshaller(self.signature, self.body)
-        body_block.marshall()
+        body_buffer = body_block._marshall()
 
         fields = []
 
@@ -287,11 +289,12 @@ class Message:
             self.message_type.value,
             self.flags.value,
             PROTOCOL_VERSION,
-            len(body_block.buffer),
+            len(body_buffer),
             self.serial,
             fields,
         ]
         header_block = Marshaller("yyyyuua(yv)", header_body)
-        header_block.marshall()
-        header_block.align(8)
-        return header_block.buffer + body_block.buffer
+        header_block._marshall()
+        header_block._align(8)
+        header_buffer = header_block._buffer()
+        return header_buffer + body_buffer