ICC Profile add utility/test

png/iccp.py

@@ -350,12 +350,6 @@ def tagblock(tag):
     return struct.pack('>L', n) + table + element
 
 
-def iccpout(out, inp):
-    """Extract ICC Profile from PNG file `inp` and write it to
-    the file `out`."""
-    out.write(png.Reader(file=inp).read()[3]['icc_profile'][1])
-
-
 def fs15f16(x):
     """Convert float to ICC s15Fixed16Number (as a Python ``int``)."""
 
@@ -522,26 +516,54 @@ def group(s, n):
     return zip(*[iter(s)] * n)
 
 
+def iccpout(out, inp, **kwargs):
+    """Extract ICC Profile from PNG file `inp` to the file `out`."""
+    out.write(png.Reader(file=inp).read()[3]['icc_profile'][1])
+
+
+def iccpadd(inp, out, addfile, **kwargs):
+    """Add ICC Profile to png file and write result to file `out`"""
+    iccfile = open(addfile, 'rb')
+    pix, meta = png.Reader(file=inp).read()[2:]
+    meta['icc_profile'] = iccfile.read()
+    iccfile.close()
+    w = png.Writer(**meta)
+    w.write(out, pix)
+
+
 def main(argv=None):
     import sys
     from getopt import getopt
+
+    def funcmode(mode):
+        """Determine function by mode"""
+        if mode == 'export':
+            return iccpout
+        elif mode == 'add':
+            return iccpadd
+
     if argv is None:
         argv = sys.argv
     argv = argv[1:]
-    opt, arg = getopt(argv, 'o:')
+    opt, arg = getopt(argv, 'o:m:a:')
+    cfg = dict(mode='export')
     if len(arg) > 0:
-        inp = open(arg[0], 'rb')
+        cfg['inp'] = open(arg[0], 'rb')
     else:
-        inp = sys.stdin
+        cfg['inp'] = sys.stdin
     for o, v in opt:
         if o == '-o':
             if v in ('-', 'stdout'):
-                f = sys.stdout
+                cfg['out'] = sys.stdout
             else:
-                f = open(v, 'wb')
-            iccpout(f, inp)
+                cfg['out'] = open(v, 'wb')
+            funcmode(cfg['mode'])(**cfg)
             if v not in ('-', 'stdout'):
-                f.close()
+                cfg['out'].close()
+        elif o == '-m':
+            cfg['mode'] = v
+        elif o == '-a':
+            cfg['addfile'] = v
 
 if __name__ == '__main__':
     main()

png/png.py

@@ -992,15 +992,15 @@ def __init__(self, width=None, height=None,
         self.planes = self.color_planes + self.alpha
 
     def set_icc_profile(self, profile=None, name='ICC Profile'):
-        if isinstance(profile, basestring):
+        if isinstance(profile, (basestring, bytes)):
             icc_profile = [name, profile]
         # TODO: more check
         else:
             icc_profile = profile
 
         if not icc_profile[0]:
             raise Error("ICC profile should have a name")
-        else:
+        elif not isinstance(icc_profile[0], bytes):
             icc_profile[0] = strtobytes(icc_profile[0])
         self.icc_profile = icc_profile
 
@@ -1157,10 +1157,14 @@ def __write_srgb(self, outfile):
                         struct.pack("!L", int(round(self.gamma * 1e5))))
         # http://www.w3.org/TR/PNG/#11iCCP
         if self.icc_profile is not None:
+            if self.compression is None or self.compression == -1:
+                comp_level = zlib.Z_DEFAULT_COMPRESSION
+            else:
+                comp_level = self.compression
             write_chunk(outfile, 'iCCP',
                         self.icc_profile[0] + zerobyte +
                         zerobyte +
-                        zlib.compress(self.icc_profile[1], self.compression))
+                        zlib.compress(self.icc_profile[1], comp_level))
 
     def __write_text(self, outfile):
         """

test_png.py

@@ -1124,6 +1124,22 @@ def testICCPexp(self):
         self.assertEqual(o.read(-1), ref.read(-1))
         ref.close()
 
+    def testICCPadd(self):
+        """Test adding ICC Profile with iccp tool"""
+        s = os.path.join(os.path.dirname(__file__),
+                         'testfiles', 'glenda.png')
+        si = os.path.join(os.path.dirname(__file__),
+                         'testfiles', 'gamma1.icc')
+        o = BytesIO()
+        _redirect_io(None, o,
+                     lambda: png.iccp.main(['iccpexp', '-madd',
+                                            '-a' + si, '-o-', s]))
+        r = png.Reader(bytes=o.getvalue())
+        meta = r.read()[3]
+        ref = open(si, 'rb')
+        o.seek(0)
+        self.assertEqual(meta['icc_profile'][1], ref.read(-1))
+        ref.close()
 
 try:
     import numpy