Make it less annoying to pickle messages

hippolyzer/lib/base/message/message.py

@@ -193,12 +193,21 @@ class Message:
         # should be set once a packet is sent / dropped to prevent accidental
         # re-sending or re-dropping
         self.finalized = False
-        # Whether message is owned by the queue or should be sent immediately
+        # Whether message is owned by a queue or should be sent immediately
         self.queued: bool = False
         self._blocks: BLOCK_DICT = {}
 
         self.add_blocks(args)
 
+    def __reduce_ex__(self, protocol):
+        reduced: Tuple[Any] = super().__reduce_ex__(protocol)
+        # https://docs.python.org/3/library/pickle.html#object.__reduce__
+        # We need to make some changes to the object state to make it serializable
+        state_dict: Dict = reduced[2][1]
+        # Have to remove the deserializer weak ref so we can pickle
+        state_dict['deserializer'] = None
+        return reduced
+
     @property
     def packet_id(self) -> Optional[int]:
         return self._packet_id

hippolyzer/lib/proxy/message_logger.py

@@ -546,7 +546,6 @@ class LLUDPMessageLogEntry(AbstractMessageLogEntry):
         # These are expensive to keep around. pickle them and un-pickle on
         # an as-needed basis.
         self._deserializer = self.message.deserializer
-        self.message.deserializer = None
         self._frozen_message = pickle.dumps(self._message, protocol=pickle.HIGHEST_PROTOCOL)
         self._message = None
 

tests/base/test_message_wrapper.py

@@ -126,8 +126,6 @@ class TestMessage(unittest.TestCase):
     def test_partial_decode_pickle(self):
         msg = self.deserial.deserialize(self.serial.serialize(self.chat_msg))
         self.assertEqual(msg.deserializer(), self.deserial)
-        # Have to remove the weak ref so we can pickle
-        msg.deserializer = None
         msg = pickle.loads(pickle.dumps(msg, protocol=pickle.HIGHEST_PROTOCOL))
 
         # We should still have the raw body at this point