From 8694f8d2ebdb639f8f2cdb9a71508ec767e19bca Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Thu, 18 May 2017 22:32:27 -0700 Subject: [PATCH] Replace magic strings colorspace and encoding with Enums --- ocrmypdf/pageinfo.py | 78 ++++++++++++++++++++++++------------------ ocrmypdf/pipeline.py | 6 ++-- tests/test_main.py | 22 ++++++------ tests/test_pageinfo.py | 7 ++-- 4 files changed, 62 insertions(+), 51 deletions(-) diff --git a/ocrmypdf/pageinfo.py b/ocrmypdf/pageinfo.py index 9d6e4e29..30117a46 100644 --- a/ocrmypdf/pageinfo.py +++ b/ocrmypdf/pageinfo.py @@ -11,48 +11,57 @@ from collections import namedtuple from collections.abc import MutableMapping, Mapping import warnings from pathlib import Path +from enum import Enum matrix_mult = pypdf.pdf.utils.matrixMultiply +Colorspace = Enum('Colorspace', + 'gray rgb cmyk lab icc index sep devn pattern jpeg2000') + +Encoding = Enum('Encoding', + 'ccitt jpeg jpeg2000 jbig2 asciihex ascii85 lzw flate ' + \ + 'runlength') + + FRIENDLY_COLORSPACE = { - '/DeviceGray': 'gray', - '/CalGray': 'gray', - '/DeviceRGB': 'rgb', - '/CalRGB': 'rgb', - '/DeviceCMYK': 'cmyk', - '/Lab': 'lab', - '/ICCBased': 'icc', - '/Indexed': 'index', - '/Separation': 'sep', - '/DeviceN': 'devn', - '/Pattern': '-', - '/G': 'gray', # Abbreviations permitted in inline images - '/RGB': 'rgb', - '/CMYK': 'cmyk', - '/I': 'index', + '/DeviceGray': Colorspace.gray, + '/CalGray': Colorspace.gray, + '/DeviceRGB': Colorspace.rgb, + '/CalRGB': Colorspace.rgb, + '/DeviceCMYK': Colorspace.cmyk, + '/Lab': Colorspace.lab, + '/ICCBased': Colorspace.icc, + '/Indexed': Colorspace.index, + '/Separation': Colorspace.sep, + '/DeviceN': Colorspace.devn, + '/Pattern': Colorspace.pattern, + '/G': Colorspace.gray, # Abbreviations permitted in inline images + '/RGB': Colorspace.rgb, + '/CMYK': Colorspace.cmyk, + '/I': Colorspace.index, } FRIENDLY_ENCODING = { - '/CCITTFaxDecode': 'ccitt', - '/DCTDecode': 'jpeg', - '/JPXDecode': 'jpx', - '/JBIG2Decode': 'jbig2', - '/CCF': 'ccitt', # Abbreviations permitted in inline images - '/DCT': 'jpeg', - '/AHx': 'asciihex', - '/A85': 'ascii85', - '/LZW': 'lzw', - '/Fl': 'flate', - '/RL': 'runlength' + '/CCITTFaxDecode': Encoding.ccitt, + '/DCTDecode': Encoding.jpeg, + '/JPXDecode': Encoding.jpeg2000, + '/JBIG2Decode': Encoding.jbig2, + '/CCF': Encoding.ccitt, # Abbreviations permitted in inline images + '/DCT': Encoding.jpeg, + '/AHx': Encoding.asciihex, + '/A85': Encoding.ascii85, + '/LZW': Encoding.lzw, + '/Fl': Encoding.flate, + '/RL': Encoding.runlength } FRIENDLY_COMP = { - 'gray': 1, - 'rgb': 3, - 'cmyk': 4, - 'lab': 3, - 'index': 1 + Colorspace.gray: 1, + Colorspace.rgb: 3, + Colorspace.cmyk: 4, + Colorspace.lab: 3, + Colorspace.index: 1 } @@ -296,7 +305,8 @@ class ImageInfo: cs = cs[0] self._color = FRIENDLY_COLORSPACE.get(cs, '-') else: - self._color = 'jpx' if self._enc == 'jpx' else '?' + self._color = FRIENDLY_COLORSPACE[Colorspace.jpeg2000] \ + if self._enc == Encoding.jpeg2000 else '?' self._comp = FRIENDLY_COMP.get(self._color, '?') @@ -304,8 +314,8 @@ class ImageInfo: # but encoding must be monochrome. This happens if a monochrome image # has an ICC profile attached. Better solution would be to examine # the ICC profile. - if self._comp == '?' and self._enc in ('ccitt', 'jbig2'): - self._comp = FRIENDLY_COMP['gray'] + if self._comp == '?' and self._enc in (Encoding.ccitt, 'jbig2'): + self._comp = FRIENDLY_COMP[Colorspace.gray] @property def name(self): diff --git a/ocrmypdf/pipeline.py b/ocrmypdf/pipeline.py index 7accb1ab..9b4794b1 100644 --- a/ocrmypdf/pipeline.py +++ b/ocrmypdf/pipeline.py @@ -22,7 +22,7 @@ from PIL import Image from ruffus import formatter, regex, Pipeline, suffix from .hocrtransform import HocrTransform -from .pageinfo import PdfInfo +from .pageinfo import PdfInfo, Encoding, Colorspace from .pdfa import generate_pdfa_ps, file_claims_pdfa from .helpers import re_symlink, is_iterable_notstr, page_number from .exec import ghostscript, tesseract, qpdf @@ -400,10 +400,10 @@ def rasterize_with_ghostscript( if all(image.comp == 1 for image in pageinfo.images): if all(image.bpc == 1 for image in pageinfo.images): device = 'pngmono' - elif all(image.bpc > 1 and image.color == 'index' + elif all(image.bpc > 1 and image.color == Colorspace.index for image in pageinfo.images): device = 'png256' - elif all(image.bpc > 1 and image.color == 'gray' + elif all(image.bpc > 1 and image.color == Colorspace.gray for image in pageinfo.images): device = 'pnggray' diff --git a/tests/test_main.py b/tests/test_main.py index bdb00f2a..8805ec2c 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -5,7 +5,7 @@ from subprocess import Popen, PIPE, check_output, check_call, DEVNULL import os import shutil import pytest -from ocrmypdf.pageinfo import PdfInfo +from ocrmypdf.pageinfo import PdfInfo, Colorspace, Encoding import PyPDF2 as pypdf from ocrmypdf.exceptions import ExitCode from ocrmypdf import leptonica @@ -630,7 +630,7 @@ def test_jbig2_passthrough(spoof_tesseract_cache, resources, outpdf): env=spoof_tesseract_cache) out_pageinfo = PdfInfo(out) - assert out_pageinfo[0].images[0].enc == 'jbig2' + assert out_pageinfo[0].images[0].enc == Encoding.jbig2 def test_stdin(spoof_tesseract_noop, ocrmypdf_exec, resources, outpdf): @@ -906,16 +906,16 @@ def test_compression_preserved(spoof_tesseract_noop, ocrmypdf_exec, pdfimage = pdfinfo[0].images[0] if input_file.endswith('.png'): - assert pdfimage.enc != 'jpeg', \ + assert pdfimage.enc != Encoding.jpeg, \ "Lossless compression changed to lossy!" elif input_file.endswith('.jpg'): - assert pdfimage.enc == 'jpeg', \ + assert pdfimage.enc == Encoding.jpeg, \ "Lossy compression changed to lossless!" if im.mode.startswith('RGB') or im.mode.startswith('BGR'): - assert pdfimage.color == 'rgb', \ + assert pdfimage.color == Colorspace.rgb, \ "Colorspace changed" elif im.mode.startswith('L'): - assert pdfimage.color == 'gray', \ + assert pdfimage.color == Colorspace.gray, \ "Colorspace changed" @@ -950,16 +950,16 @@ def test_compression_changed(spoof_tesseract_noop, ocrmypdf_exec, pdfimage = pdfinfo[0].images[0] - if compression == 'jpeg': - assert pdfimage.enc == 'jpeg' + if compression == "jpeg": + assert pdfimage.enc == Encoding.jpeg elif compression == 'lossless': - assert pdfimage.enc == 'image' + assert pdfimage.enc not in (Encoding.jpeg, Encoding.jpeg2000) if im.mode.startswith('RGB') or im.mode.startswith('BGR'): - assert pdfimage.color == 'rgb', \ + assert pdfimage.color == Colorspace.rgb, \ "Colorspace changed" elif im.mode.startswith('L'): - assert pdfimage.color == 'gray', \ + assert pdfimage.color == Colorspace.gray, \ "Colorspace changed" diff --git a/tests/test_pageinfo.py b/tests/test_pageinfo.py index 663380e5..66faa548 100644 --- a/tests/test_pageinfo.py +++ b/tests/test_pageinfo.py @@ -6,6 +6,7 @@ from reportlab.pdfgen.canvas import Canvas from PIL import Image from tempfile import NamedTemporaryFile from math import isclose +from ocrmypdf.pageinfo import Colorspace, Encoding from contextlib import suppress import os import shutil @@ -64,7 +65,7 @@ def test_single_page_image(outdir): pdfimage = page.images[0] assert pdfimage.width == 8 - assert pdfimage.color == 'gray' + assert pdfimage.color == Colorspace.gray # DPI in a 1"x1" is the image width assert isclose(pdfimage.xres, 8) @@ -88,7 +89,7 @@ def test_single_page_inline_image(outdir): print(pdfinfo) pdfimage = pdfinfo[0].images[0] assert isclose(pdfimage.xres, 8) - assert pdfimage.color != '-' + assert pdfimage.color == Colorspace.rgb # reportlab produces color image assert pdfimage.width == 8 @@ -98,7 +99,7 @@ def test_jpeg(resources, outdir): pdfinfo = pageinfo.PdfInfo(filename) pdfimage = pdfinfo[0].images[0] - assert pdfimage.enc == 'jpeg' + assert pdfimage.enc == Encoding.jpeg assert isclose(pdfimage.xres, 150)