Replace magic strings colorspace and encoding with Enums

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):

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'
 

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"
 
 

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)