Add support for saving CMYK images

drawBot/context/imageContext.py

@@ -32,6 +32,12 @@ def _tiffCompressionConverter(value):
         return t.get(value.lower(), AppKit.NSTIFFCompressionNone)
 
 
+def _colorSpaceConverter(value):
+    if value == "CMYK":
+        return AppKit.NSDeviceCMYKColorSpace
+    return AppKit.NSCalibratedRGBColorSpace
+
+
 _nsImageOptions = {
     # DrawBot Key: (
     #   NSImage property key,
@@ -118,6 +124,7 @@ class ImageContext(PDFContext):
             "A Boolean value that specifies whether subpixel quantization of glyphs is allowed. Default is True.",
         ),
         ("multipage", "Output a numbered image for each page or frame in the document."),
+        ("colorSpace", "Set color space (RGB, CMYK). Default is RGB."),
     ]
 
     ensureEvenPixelDimensions = False
@@ -138,6 +145,10 @@ def _writeDataToFile(self, data, path, options):
         imageResolution = options.get("imageResolution", 72.0)
         antiAliasing = options.get("antiAliasing", True)
         fontSubpixelQuantization = options.get("fontSubpixelQuantization", True)
+
+        colorSpace = options.get("colorSpace", "RGB")
+        colorSpaceName = _colorSpaceConverter(colorSpace)
+
         properties = {}
         for key, value in options.items():
             if key in _nsImageOptions:
@@ -154,6 +165,7 @@ def _writeDataToFile(self, data, path, options):
                     antiAliasing=antiAliasing,
                     fontSubpixelQuantization=fontSubpixelQuantization,
                     imageResolution=imageResolution,
+                    colorSpaceName=colorSpaceName,
                 )
                 if self.ensureEvenPixelDimensions:
                     if imageRep.pixelsWide() % 2 or imageRep.pixelsHigh() % 2:
@@ -189,13 +201,17 @@ def _makeBitmapImageRep(
     elif nsImage is not None:
         width, height = nsImage.size()
 
+    hasAlpha = True
+    if colorSpaceName == AppKit.NSDeviceCMYKColorSpace:
+        hasAlpha = False  # Quartz doesn’t support alpha for CMYK bitmaps.
+
     rep = AppKit.NSBitmapImageRep.alloc().initWithBitmapDataPlanes_pixelsWide_pixelsHigh_bitsPerSample_samplesPerPixel_hasAlpha_isPlanar_colorSpaceName_bytesPerRow_bitsPerPixel_(
         None,  # planes
         int(width * scaleFactor),  # pixelsWide
         int(height * scaleFactor),  # pixelsHigh
         8,  # bitsPerSample
         4,  # samplesPerPixel
-        True,  # hasAlpha
+        hasAlpha,  # hasAlpha
         False,  # isPlanar
         colorSpaceName,  # colorSpaceName
         0,  # bytesPerRow

tests/testExport.py

@@ -1,11 +1,13 @@
 import glob
+import io
 import os
 import random
 import sys
 import unittest
 
 import AppKit  # type: ignore
 import PIL
+from PIL import ImageCms
 from testSupport import (
     DrawBotBaseTest,
     StdOutCollector,
@@ -444,6 +446,55 @@ def test_formattedStringURL_svg(self):
             readData(path), readData(expectedPath), "Files %r and %s are not the same" % (path, expectedPath)
         )
 
+    def test_export_color_space_cmyk(self):
+        drawBot.newDrawing()
+        drawBot.size(1000, 1000)
+        drawBot.cmykFill(1, 0, 0, 0)
+        drawBot.rect(0, 0, 250, 1000)
+        drawBot.cmykFill(0, 1, 0, 0)
+        drawBot.rect(250, 0, 250, 1000)
+        drawBot.cmykFill(0, 0, 1, 0)
+        drawBot.rect(500, 0, 250, 1000)
+        drawBot.cmykFill(0, 0, 0, 1)
+        drawBot.rect(750, 0, 250, 1000)
+
+        rgba_path = os.path.join(tempTestDataDir, "cmyk_export_default.tiff")
+        cmyk_path = os.path.join(tempTestDataDir, "cmyk_export_with_color_space.tiff")
+        cmyk_with_profile_path = os.path.join(tempTestDataDir, "cmyk_export_with_color_space_and_profile.tiff")
+
+        drawBot.saveImage(rgba_path)
+        drawBot.saveImage(cmyk_path, colorSpace="CMYK")
+
+        icc_data = AppKit.NSData.dataWithContentsOfFile_("/System/Library/ColorSync/Profiles/Generic CMYK Profile.icc")
+        drawBot.saveImage(cmyk_with_profile_path, colorSpace="CMYK", imageColorSyncProfileData=icc_data)
+
+        drawBot.endDrawing()
+
+        rgba_image = PIL.Image.open(rgba_path)
+        self.assertEqual(rgba_image.mode, "RGBA")
+
+        rgba_image_profile = ImageCms.ImageCmsProfile(
+            io.BytesIO(rgba_image.info.get("icc_profile"))
+        )
+        self.assertEqual(
+            rgba_image_profile.profile.profile_description,
+            "Generic RGB Profile"
+        )
+
+        cmyk_image = PIL.Image.open(cmyk_path)
+        self.assertEqual(cmyk_image.mode, "CMYK")
+        self.assertIsNone(cmyk_image.info.get("icc_profile"))
+
+        cmyk_with_profile_image = PIL.Image.open(cmyk_with_profile_path)
+        self.assertEqual(cmyk_with_profile_image.mode, "CMYK")
+        cmyk_with_profile_image_profile = ImageCms.ImageCmsProfile(
+            io.BytesIO(cmyk_with_profile_image.info.get("icc_profile"))
+        )
+        self.assertEqual(
+            cmyk_with_profile_image_profile.profile.profile_description,
+            "Generic CMYK Profile"
+        )
+
 
 if __name__ == "__main__":
     import doctest