#!/usr/bin/env python
# $URL: https://pypng.googlecode.com/svn/trunk/code/pngchunk $
# $Rev: 107 $
# pngchunk
# Chunk editing/extraction tool.

import struct
import warnings

# Local module.
import png

def chunk(out, inp, l):
    """Process the input PNG file to the output, chunk by chunk.  Chunks
    can be inserted, removed, replaced, or sometimes edited.  Generally, 
    chunks are not inspected, so pixel data (in the ``IDAT`` chunks)
    cannot be modified.  `l` should be a list of (*chunktype*,
    *content*) pairs.  *chunktype* is usually the type of the PNG chunk,
    specified as a 4-byte Python string, and *content* is the chunk's
    content, also as a string.

    This function *knows* about certain chunk types and will
    automatically convert from Python friendly representations to
    string-of-bytes.

    chunktype
    'gamma'     'gAMA'  float
    'sigbit'    'sBIT'  int, or tuple of length 1,2 or 3

    Note that the length of the strings used to identify *friendly*
    chunk types is greater than 4, hence they cannot be confused with
    canonical chunk types.

    Chunk types, if specified using the 4-byte syntax, need not be
    official PNG chunks at all.  Non-standard chunks can be created.
    """

    def canonical(p):
        """Take a pair (*chunktype*, *content*), and return canonical
        representation (*chunktype*, *content*) where `chunktype` is the
        4-byte PNG chunk type and `content` is a string.
        """

        t,v = p
        if len(t) == 4:
            return t,v
        if t == 'gamma':
            t = 'gAMA'
            v = int(round(1e5*v))
            v = struct.pack('>I', v)
        elif t == 'sigbit':
            t = 'sBIT'
            try:
                v[0]
            except TypeError:
                v = (v,)
            v = struct.pack('%dB' % len(v), *v)
        else:
            warnings.warn('Unknown chunk type %r' % t)
        return t[:4],v

    l = map(canonical, l)
    chunktypes,_ = zip(*l)
    chunktypes = list(chunktypes)
    # Some chunks automagically replace ones that are present in the
    # source PNG.  There can only be one of each of these chunk types.
    # Create a 'replace' dictionary to record these chunks.
    d = dict(l)
    replacing = ['gAMA', 'sBIT', 'PLTE', 'tRNS', 'sPLT', 'IHDR']
    replace = dict()
    for t in replacing:
        if t in d:
            replace[t] = d[t]
            del d[t]
            i = chunktypes.index(t)
            del l[i]
    r = png.Reader(file=inp)
    chunks = r.chunks()
    def iterchunks():
        for t,v in chunks:
            if t in replace:
                yield t,replace[t]
                del replace[t]
                continue
            if t == 'IDAT' and replace:
                # Insert into the output any chunks that are on the
                # replace list.  We haven't output them yet, because we
                # didn't see an original chunk of the same type to
                # replace.  Thus the "replace" is actually an "insert".
                for u,w in replace.items():
                    yield u,w
                    del replace[u]
            if t == 'IDAT' and l:
                for item in l:
                    yield item
                del l[:]
            yield t,v
    return png.write_chunks(out, iterchunks())


def main(argv=None):
    import sys
    from getopt import getopt

    if argv is None:
        argv = sys.argv

    argv = argv[1:]
    opt,arg = getopt(argv, '', ['gamma=', 'sigbit='])
    k = []
    for o,v in opt:
        if o in ['--gamma']:
            k.append(('gamma', float(v)))
        if o in ['--sigbit']:
            k.append(('sigbit', int(v)))
    if len(arg) > 0:
        f = open(arg[0], 'rb')
    else:
        f = sys.stdin
    return chunk(sys.stdout, f, k)


if __name__ == '__main__':
    main()
