From 18d637ea340e46a06806846d586a459c0334b7fd Mon Sep 17 00:00:00 2001
From: Stephan Richter <stephan.richter@gmail.com>
Date: Mon, 22 May 2017 12:31:52 -0400
Subject: [PATCH] Make sure we are using the correct IO class in Python 2.

---
 src/tests/__init__.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/tests/__init__.py b/src/tests/__init__.py
index 47d84ed..caa201b 100644
--- a/src/tests/__init__.py
+++ b/src/tests/__init__.py
@@ -1,14 +1,23 @@
 import unittest
 
-import os
 import img2pdf
+import os
+import struct
+import sys
 import zlib
 from PIL import Image
-from io import BytesIO
-import struct
+from io import StringIO, BytesIO
 
 HERE = os.path.dirname(__file__)
 
+PY3 = sys.version_info[0] >= 3
+
+if PY3:
+    PdfReaderIO = StringIO
+else:
+    PdfReaderIO = BytesIO
+
+
 # convert +set date:create +set date:modify -define png:exclude-chunk=time
 
 # we define some variables so that the table below can be narrower
@@ -470,10 +479,9 @@ def test_suite():
                 orig_imgdata = inf.read()
             output = img2pdf.convert(orig_imgdata, nodate=True,
                                      with_pdfrw=with_pdfrw)
-            from io import StringIO, BytesIO
             from pdfrw import PdfReader, PdfName, PdfWriter
             from pdfrw.py23_diffs import convert_load, convert_store
-            x = PdfReader(StringIO(convert_load(output)))
+            x = PdfReader(PdfReaderIO(convert_load(output)))
             self.assertEqual(sorted(x.keys()), [PdfName.Info, PdfName.Root,
                              PdfName.Size])
             self.assertIn(x.Root.Pages.Count, ('1', '2'))