src/lib/png.py 4.6 K raw
1
"""Minimal RGBA / indexed PNG reader and writer.
2
3
Only what the theme's assets need: 8-bit truecolour-alpha (type 6) and indexed
4
(type 3) images, non-interlaced. Written against zlib from the stdlib so the
5
theme can be rebuilt without Pillow or ImageMagick installed.
6
"""
7
8
import struct
9
import zlib
10
11
12
def _chunks(data):
13
    assert data[:8] == b"\x89PNG\r\n\x1a\n", "not a PNG"
14
    i = 8
15
    out = []
16
    while i < len(data):
17
        (ln,) = struct.unpack(">I", data[i:i + 4])
18
        out.append([data[i + 4:i + 8], data[i + 8:i + 8 + ln]])
19
        i += 8 + ln + 4
20
    return out
21
22
23
def _pack(chunks):
24
    out = bytearray(b"\x89PNG\r\n\x1a\n")
25
    for typ, body in chunks:
26
        out += struct.pack(">I", len(body)) + typ + body
27
        out += struct.pack(">I", zlib.crc32(typ + body) & 0xFFFFFFFF)
28
    return bytes(out)
29
30
31
def _paeth(a, b, c):
32
    p = a + b - c
33
    pa, pb, pc = abs(p - a), abs(p - b), abs(p - c)
34
    if pa <= pb and pa <= pc:
35
        return a
36
    return b if pb <= pc else c
37
38
39
def _unfilter(raw, width, height, bpp):
40
    """Reverse the per-scanline filters, returning flat sample bytes."""
41
    stride = width * bpp
42
    out = bytearray(stride * height)
43
    pos = 0
44
    for y in range(height):
45
        ft = raw[pos]
46
        pos += 1
47
        line = bytearray(raw[pos:pos + stride])
48
        pos += stride
49
        base = y * stride
50
        prev = base - stride
51
        if ft == 0:
52
            pass
53
        elif ft == 1:
54
            for x in range(bpp, stride):
55
                line[x] = (line[x] + line[x - bpp]) & 0xFF
56
        elif ft == 2:
57
            if y:
58
                for x in range(stride):
59
                    line[x] = (line[x] + out[prev + x]) & 0xFF
60
        elif ft == 3:
61
            for x in range(stride):
62
                a = line[x - bpp] if x >= bpp else 0
63
                b = out[prev + x] if y else 0
64
                line[x] = (line[x] + ((a + b) >> 1)) & 0xFF
65
        elif ft == 4:
66
            for x in range(stride):
67
                a = line[x - bpp] if x >= bpp else 0
68
                b = out[prev + x] if y else 0
69
                c = out[prev + x - bpp] if (y and x >= bpp) else 0
70
                line[x] = (line[x] + _paeth(a, b, c)) & 0xFF
71
        else:
72
            raise ValueError("unknown filter %d" % ft)
73
        out[base:base + stride] = line
74
    return out
75
76
77
def _filter_up(flat, width, height, bpp):
78
    """Re-encode with the Up filter; cheap and compresses these assets well."""
79
    stride = width * bpp
80
    out = bytearray()
81
    for y in range(height):
82
        base = y * stride
83
        line = flat[base:base + stride]
84
        if y == 0:
85
            out.append(0)
86
            out += line
87
        else:
88
            prev = flat[base - stride:base]
89
            out.append(2)
90
            out += bytes((line[x] - prev[x]) & 0xFF for x in range(stride))
91
    return bytes(out)
92
93
94
class Png:
95
    def __init__(self, path):
96
        self.path = path
97
        self.chunks = _chunks(open(path, "rb").read())
98
        ihdr = dict((t, b) for t, b in self.chunks)[b"IHDR"]
99
        (self.width, self.height, self.depth, self.color_type,
100
         _, _, self.interlace) = struct.unpack(">IIBBBBB", ihdr)
101
        assert self.interlace == 0, "interlaced PNGs unsupported"
102
103
    def _get(self, typ):
104
        for t, b in self.chunks:
105
            if t == typ:
106
                return b
107
        return None
108
109
    def _set(self, typ, body):
110
        for c in self.chunks:
111
            if c[0] == typ:
112
                c[1] = body
113
                return
114
        raise KeyError(typ)
115
116
    # -- indexed ------------------------------------------------------------
117
    def palette(self):
118
        p = self._get(b"PLTE")
119
        return [tuple(p[i:i + 3]) for i in range(0, len(p), 3)]
120
121
    def set_palette(self, pal):
122
        self._set(b"PLTE", b"".join(bytes(c) for c in pal))
123
124
    # -- truecolour+alpha ---------------------------------------------------
125
    def rgba(self):
126
        assert self.color_type == 6 and self.depth == 8, "expected 8-bit RGBA"
127
        idat = b"".join(b for t, b in self.chunks if t == b"IDAT")
128
        flat = _unfilter(zlib.decompress(idat), self.width, self.height, 4)
129
        return [tuple(flat[i:i + 4]) for i in range(0, len(flat), 4)]
130
131
    def set_rgba(self, pixels):
132
        flat = bytearray()
133
        for px in pixels:
134
            flat += bytes(px)
135
        raw = _filter_up(flat, self.width, self.height, 4)
136
        body = zlib.compress(raw, 9)
137
        # collapse any split IDATs into one
138
        self.chunks = [c for c in self.chunks if c[0] != b"IDAT"]
139
        end = next(i for i, c in enumerate(self.chunks) if c[0] == b"IEND")
140
        self.chunks.insert(end, [b"IDAT", body])
141
142
    def save(self, path=None):
143
        open(path or self.path, "wb").write(_pack(self.chunks))