From 43d650e78c47ea67a42442ad54d459ba1da8c823 Mon Sep 17 00:00:00 2001
From: "James R. Barlow" <jim@purplerock.ca>
Date: Sat, 25 Apr 2020 03:50:11 -0700
Subject: [PATCH] Fix issue where only first PNG-style image would be optimized

---
 src/ocrmypdf/optimize.py |  4 ++--
 tests/test_optimize.py   | 42 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/src/ocrmypdf/optimize.py b/src/ocrmypdf/optimize.py
index 839a8ff9..99211698 100644
--- a/src/ocrmypdf/optimize.py
+++ b/src/ocrmypdf/optimize.py
@@ -421,9 +421,9 @@ def transcode_pngs(pike, images, image_name_fn, root, log, options):
             )
             continue
         if compdata.type == leptonica.lept.L_FLATE_ENCODE:
-            return rewrite_png(pike, im_obj, compdata, log)
+            rewrite_png(pike, im_obj, compdata, log)
         elif compdata.type == leptonica.lept.L_G4_ENCODE:
-            return rewrite_png_as_g4(pike, im_obj, compdata, log)
+            rewrite_png_as_g4(pike, im_obj, compdata, log)
 
 
 def rewrite_png_as_g4(pike, im_obj, compdata, log):
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index 0c78d653..b9368849 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -18,10 +18,12 @@
 import logging
 from os import fspath
 from pathlib import Path
+from unittest.mock import patch
 
+import img2pdf
 import pikepdf
 import pytest
-from PIL import Image
+from PIL import Image, ImageDraw
 
 from ocrmypdf import optimize as opt
 from ocrmypdf.exec import jbig2enc, pngquant
@@ -130,3 +132,41 @@ def test_flate_to_jbig2(resources, outdir, spoof_tesseract_noop):
     pdf = pikepdf.open(outdir / 'out.pdf')
     pim = pikepdf.PdfImage(next(iter(pdf.pages[0].images.values())))
     assert pim.filters[0] == '/JBIG2Decode'
+
+
+def test_multiple_pngs(resources, outdir, spoof_tesseract_noop):
+    with Path.open(outdir / 'in.pdf', 'wb') as inpdf:
+        img2pdf.convert(
+            fspath(resources / 'baiona_colormapped.png'),
+            fspath(resources / 'baiona_gray.png'),
+            with_pdfrw=False,
+            outputstream=inpdf,
+        )
+
+    def mockquant(input_file, output_file, _quality_min, _quality_max):
+        with Image.open(input_file) as im:
+            draw = ImageDraw.Draw(im)
+            draw.rectangle((0, 0, im.width, im.height), fill=128)
+            im.save(output_file)
+
+    with patch('ocrmypdf.optimize.pngquant.quantize', new=mockquant):
+        check_ocrmypdf(
+            outdir / 'in.pdf',
+            outdir / 'out.pdf',
+            '--optimize',
+            '3',
+            '--jobs',
+            '1',
+            '--use-threads',
+            '--output-type',
+            'pdf',
+            env=spoof_tesseract_noop,
+        )
+
+    with pikepdf.open(outdir / 'in.pdf') as inpdf, pikepdf.open(
+        outdir / 'out.pdf'
+    ) as outpdf:
+        for n in range(len(inpdf.pages)):
+            inim = next(iter(inpdf.pages[n].images.values()))
+            outim = next(iter(outpdf.pages[n].images.values()))
+            assert len(outim.read_raw_bytes()) < len(inim.read_raw_bytes()), n