fix: avoid writing empty image files (#139)

* fix: avoid writing empty files

* fix: None

Author: dhdaines

playa/cli.py

@@ -87,6 +87,7 @@
 
 import playa
 from playa import Document, Page, PDFPasswordIncorrect, asobj
+from playa.color import ColorSpace
 from playa.data.content import Image
 from playa.data.metadata import asobj_document
 from playa.outline import Outline
@@ -517,30 +518,50 @@ def get_one_image(stream: ContentStream, path: Path) -> Path:
     else:
         # Otherwise, try to write a PNM file
         bits = stream.bits
-        colorspace = stream.colorspace
-        ncomponents = colorspace.ncomponents
-        if colorspace.name == "Indexed":
-            from playa.color import get_colorspace
-
-            assert isinstance(colorspace.spec, list)
-            _, underlying, _, _ = colorspace.spec
-            underlying = get_colorspace(resolve1(underlying))
-            if underlying is not None:
-                ncomponents = underlying.ncomponents
+        colorspace: Union[ColorSpace, None] = stream.colorspace
+        ncomponents = 1
+        if colorspace is not None:
+            ncomponents = colorspace.ncomponents
+            if colorspace.name == "Indexed":
+                from playa.color import get_colorspace
+
+                assert isinstance(colorspace.spec, list)
+                _, underlying, _, _ = colorspace.spec
+                colorspace = get_colorspace(resolve1(underlying))
+                if colorspace is not None:
+                    ncomponents = colorspace.ncomponents
         if bits == 1:
             path = path.with_suffix(".pbm")
+        elif colorspace is None or colorspace.name not in ("DeviceGray", "DeviceRGB"):
+            path = path.with_suffix(".dat")
+            LOG.warning(
+                "Unsupported colorspace %s, writing data to %s", asobj(colorspace), path
+            )
         elif ncomponents == 1:
             path = path.with_suffix(".pgm")
         elif ncomponents == 3:
             path = path.with_suffix(".ppm")
-        try:
+        elif ncomponents == 3:
+            path = path.with_suffix(".ppm")
+        else:
+            path = path.with_suffix(".dat")
+            LOG.warning(
+                "Unsupported colorspace %s, writing data to %s", asobj(colorspace), path
+            )
+
+        if path.suffix != ".dat":
+            try:
+                with open(path, "wb") as outfh:
+                    stream.write_pnm(outfh)
+            except ValueError:
+                datpath = path.with_suffix(".dat")
+                LOG.exception(
+                    "Failed to write PNM to %s, writing data to %s", path, datpath
+                )
+                path = datpath
+
+        if path.suffix == ".dat":
             with open(path, "wb") as outfh:
-                stream.write_pnm(outfh)
-        except ValueError:
-            # Fall back to a binary file
-            datpath = path.with_suffix(".dat")
-            LOG.warning("Failed to write PNM to %s, writing data to %s", path, datpath)
-            with open(datpath, "wb") as outfh:
                 outfh.write(stream.buffer)
     return path
 

playa/pdftypes.py

@@ -632,6 +632,7 @@ def write_pnm(self, outfh: BinaryIO) -> None:
         bits = self.bits
         colorspace = self.colorspace
         data = self.buffer
+        ftype: Union[bytes, None] = None
         if bits == 1:
             ftype = b"P4"
         elif colorspace.name == "DeviceGray":
@@ -662,7 +663,7 @@ def write_pnm(self, outfh: BinaryIO) -> None:
                 for i in unpack_indexed_image_data(data, bits, self.width, self.height)
                 for b in lookup[channels * i : channels * (i + 1)]
             )
-        else:
+        if ftype is None:
             raise ValueError("Unsupported colorspace: %r" % (self.colorspace,))
         max_value = (1 << bits) - 1
         outfh.write(b"%s %d %d\n" % (ftype, self.width, self.height))

tests/test_lazy_api.py

@@ -2,7 +2,6 @@
 Test the ContentObject API for pages.
 """
 
-import sys
 from pathlib import Path
 
 import playa