From da787214e1d196992aecf269c9d0105e4c934a4d Mon Sep 17 00:00:00 2001
From: Dana Powers <dana.powers@rd.io>
Date: Sat, 23 Jan 2016 15:25:05 -0800
Subject: [PATCH] Support encode and repr on raw BytesIO MessageSets (used in
 new producer)

---
 kafka/protocol/message.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/kafka/protocol/message.py b/kafka/protocol/message.py
index dffb1bb..fb54049 100644
--- a/kafka/protocol/message.py
+++ b/kafka/protocol/message.py
@@ -88,6 +88,13 @@ class MessageSet(AbstractType):
 
     @classmethod
     def encode(cls, items, size=True, recalc_message_size=True):
+        # RecordAccumulator encodes messagesets internally
+        if isinstance(items, io.BytesIO):
+            size = Int32.decode(items)
+            # rewind and return all the bytes
+            items.seek(-4, 1)
+            return items.read(size + 4)
+
         encoded_values = []
         for (offset, message_size, message) in items:
             if isinstance(message, Message):
@@ -143,4 +150,9 @@ class MessageSet(AbstractType):
 
     @classmethod
     def repr(cls, messages):
+        if isinstance(messages, io.BytesIO):
+            offset = messages.tell()
+            decoded = cls.decode(messages)
+            messages.seek(offset)
+            messages = decoded
         return '[' + ', '.join([cls.ITEM.repr(m) for m in messages]) + ']'