Make all unit tests pass on py3.3/3.4

1 files changed, 11 insertions, 8 deletions

diff --git a/kafka/codec.py b/kafka/codec.py
index 206ddb4..2279200 100644
--- a/kafka/codec.py
+++ b/kafka/codec.py
@@ -1,8 +1,11 @@
-from cStringIO import StringIO
+from io import BytesIO
 import gzip
 import struct
 
-_XERIAL_V1_HEADER = (-126, 'S', 'N', 'A', 'P', 'P', 'Y', 0, 1, 1)
+import six
+from six.moves import xrange
+
+_XERIAL_V1_HEADER = (-126, b'S', b'N', b'A', b'P', b'P', b'Y', 0, 1, 1)
 _XERIAL_V1_FORMAT = 'bccccccBii'
 
 try:
@@ -21,7 +24,7 @@ def has_snappy():
 
 
 def gzip_encode(payload):
-    buffer = StringIO()
+    buffer = BytesIO()
     handle = gzip.GzipFile(fileobj=buffer, mode="w")
     handle.write(payload)
     handle.close()
@@ -32,7 +35,7 @@ def gzip_encode(payload):
 
 
 def gzip_decode(payload):
-    buffer = StringIO(payload)
+    buffer = BytesIO(payload)
     handle = gzip.GzipFile(fileobj=buffer, mode='r')
     result = handle.read()
     handle.close()
@@ -68,9 +71,9 @@ def snappy_encode(payload, xerial_compatible=False, xerial_blocksize=32 * 1024):
             for i in xrange(0, len(payload), xerial_blocksize):
                 yield payload[i:i+xerial_blocksize]
 
-        out = StringIO()
+        out = BytesIO()
 
-        header = ''.join([struct.pack('!' + fmt, dat) for fmt, dat
+        header = b''.join([struct.pack('!' + fmt, dat) for fmt, dat
             in zip(_XERIAL_V1_FORMAT, _XERIAL_V1_HEADER)])
 
         out.write(header)
@@ -121,8 +124,8 @@ def snappy_decode(payload):
 
     if _detect_xerial_stream(payload):
         # TODO ? Should become a fileobj ?
-        out = StringIO()
-        byt = buffer(payload[16:])
+        out = BytesIO()
+        byt = payload[16:]
         length = len(byt)
         cursor = 0