From 79d52a50bd9651fa489e81935bda04d53285b2c1 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Fri, 4 Nov 2022 00:10:29 +0100
Subject: [PATCH] feat: optimize unmarshaller by dropping exception that was
 only used internally (#145)

---
 src/dbus_fast/_private/unmarshaller.pxd |  3 +--
 src/dbus_fast/_private/unmarshaller.py  | 34 ++++++++-----------------
 2 files changed, 11 insertions(+), 26 deletions(-)

diff --git a/src/dbus_fast/_private/unmarshaller.pxd b/src/dbus_fast/_private/unmarshaller.pxd
index ee96541..43a0d08 100644
--- a/src/dbus_fast/_private/unmarshaller.pxd
+++ b/src/dbus_fast/_private/unmarshaller.pxd
@@ -71,6 +71,7 @@ cdef unsigned int TOKEN_O_AS_INT
 cdef unsigned int TOKEN_S_AS_INT
 cdef unsigned int TOKEN_G_AS_INT
 
+cdef object MARSHALL_STREAM_END_ERROR
 
 cpdef get_signature_tree
 
@@ -87,8 +88,6 @@ cdef inline unsigned short _cast_uint16_native(const char *  payload, unsigned i
     return u16p[0]
 
 
-cdef class MarshallerStreamEndError(Exception):
-    pass
 
 cdef class Unmarshaller:
 
diff --git a/src/dbus_fast/_private/unmarshaller.py b/src/dbus_fast/_private/unmarshaller.py
index 39074bf..06b0535 100644
--- a/src/dbus_fast/_private/unmarshaller.py
+++ b/src/dbus_fast/_private/unmarshaller.py
@@ -109,6 +109,8 @@ HEADER_MESSAGE_ARG_NAME = {
 
 READER_TYPE = Callable[["Unmarshaller", SignatureType], Any]
 
+MARSHALL_STREAM_END_ERROR = BlockingIOError
+
 
 def unpack_parser_factory(unpack_from: Callable, size: int) -> READER_TYPE:
     """Build a parser that unpacks the bytes using the given unpack_from function."""
@@ -134,17 +136,6 @@ def build_simple_parsers(
     return parsers
 
 
-class MarshallerStreamEndError(Exception):
-    """This exception is raised when the end of the stream is reached.
-
-    This means more data is expected on the wire that has not yet been
-    received. The caller should call unmarshall later when more data is
-    available.
-    """
-
-    pass
-
-
 try:
     import cython
 except ImportError:
@@ -235,14 +226,9 @@ class Unmarshaller:
     def _read_sock(self, length: int) -> bytes:
         """reads from the socket, storing any fds sent and handling errors
         from the read itself"""
-
-        try:
-            msg, ancdata, _flags, _addr = self._sock.recvmsg(
-                length, UNIX_FDS_CMSG_LENGTH
-            )
-        except BlockingIOError:
-            raise MarshallerStreamEndError()
-
+        # This will raise BlockingIOError if there is no data to read
+        # which we store in the MARSHALL_STREAM_END_ERROR object
+        msg, ancdata, _flags, _addr = self._sock.recvmsg(length, UNIX_FDS_CMSG_LENGTH)
         for level, type_, data in ancdata:
             if not (level == SOL_SOCKET and type_ == SCM_RIGHTS):
                 continue
@@ -256,7 +242,7 @@ class Unmarshaller:
         """
         Read from underlying socket into buffer.
 
-        Raises MarshallerStreamEndError if there is not enough data to be read.
+        Raises BlockingIOError if there is not enough data to be read.
 
         :arg pos:
             The pos to read to. If not enough bytes are available in the
@@ -274,10 +260,10 @@ class Unmarshaller:
         if data == b"":
             raise EOFError()
         if data is None:
-            raise MarshallerStreamEndError()
+            raise MARSHALL_STREAM_END_ERROR
         self._buf += data
         if len(data) + start_len != pos:
-            raise MarshallerStreamEndError()
+            raise MARSHALL_STREAM_END_ERROR
 
     def read_uint32_unpack(self, type_: SignatureType) -> int:
         return self._read_uint32_unpack()
@@ -610,7 +596,7 @@ class Unmarshaller:
     def unmarshall(self) -> Optional[Message]:
         """Unmarshall the message.
 
-        The underlying read function will raise MarshallerStreamEndError
+        The underlying read function will raise BlockingIOError if the
         if there are not enough bytes in the buffer. This allows unmarshall
         to be resumed when more data comes in over the wire.
         """
@@ -618,7 +604,7 @@ class Unmarshaller:
             if not self._msg_len:
                 self._read_header()
             self._read_body()
-        except MarshallerStreamEndError:
+        except MARSHALL_STREAM_END_ERROR:
             return None
         return self._message