From d931f02709ca70e6d68636c49ee3d5e2f778c2e9 Mon Sep 17 00:00:00 2001
From: Johannes 'josch' Schauer <josch@mister-muffin.de>
Date: Wed, 18 Jul 2018 00:05:15 +0200
Subject: [PATCH] src/tests/__init__.py: handle PNG input images properly

---
 src/tests/__init__.py | 55 +++++++++++++++++++++++++------------------
 1 file changed, 32 insertions(+), 23 deletions(-)

diff --git a/src/tests/__init__.py b/src/tests/__init__.py
index 98ebeaf..77b8148 100644
--- a/src/tests/__init__.py
+++ b/src/tests/__init__.py
@@ -589,30 +589,39 @@ def test_suite():
                     # to the pixel data of the input image
                     imgdata = zlib.decompress(
                         convert_store(cur_page.Resources.XObject.Im0.stream))
-                    colorspace = imgprops.ColorSpace
-                    if colorspace == PdfName.DeviceGray:
-                        colorspace = 'L'
-                    elif colorspace == PdfName.DeviceRGB:
-                        colorspace = 'RGB'
-                    elif colorspace == PdfName.DeviceCMYK:
-                        colorspace = 'CMYK'
+                    if imgprops.DecodeParms:
+                        if orig_img.format == 'PNG':
+                            pngidat, palette = img2pdf.parse_png(orig_imgdata)
+                        else:
+                            pngbuffer = BytesIO()
+                            orig_img.save(pngbuffer, format="png")
+                            pngidat, palette = img2pdf.parse_png(pngbuffer.getvalue())
+                        self.assertEqual(zlib.decompress(pngidat), imgdata)
                     else:
-                        raise Exception("invalid colorspace")
-                    im = Image.frombytes(colorspace, (int(imgprops.Width),
-                                                      int(imgprops.Height)),
-                                         imgdata)
-                    if orig_img.mode == '1':
-                        self.assertEqual(im.tobytes(),
-                                         orig_img.convert("L").tobytes())
-                    elif orig_img.mode not in ("RGB", "L", "CMYK", "CMYK;I"):
-                        self.assertEqual(im.tobytes(),
-                                         orig_img.convert("RGB").tobytes())
-                    # the python-pil version 2.3.0-1ubuntu3 in Ubuntu does not
-                    # have the close() method
-                    try:
-                        im.close()
-                    except AttributeError:
-                        pass
+                        colorspace = imgprops.ColorSpace
+                        if colorspace == PdfName.DeviceGray:
+                            colorspace = 'L'
+                        elif colorspace == PdfName.DeviceRGB:
+                            colorspace = 'RGB'
+                        elif colorspace == PdfName.DeviceCMYK:
+                            colorspace = 'CMYK'
+                        else:
+                            raise Exception("invalid colorspace")
+                        im = Image.frombytes(colorspace, (int(imgprops.Width),
+                                                          int(imgprops.Height)),
+                                             imgdata)
+                        if orig_img.mode == '1':
+                            self.assertEqual(im.tobytes(),
+                                             orig_img.convert("L").tobytes())
+                        elif orig_img.mode not in ("RGB", "L", "CMYK", "CMYK;I"):
+                            self.assertEqual(im.tobytes(),
+                                             orig_img.convert("RGB").tobytes())
+                        # the python-pil version 2.3.0-1ubuntu3 in Ubuntu does not
+                        # have the close() method
+                        try:
+                            im.close()
+                        except AttributeError:
+                            pass
             # now use pdfrw to parse and then write out both pdfs and check the
             # result for equality
             y = PdfReader(out)