Generating a PNG in Python

png_out.py

import base64, io, struct, zlib
from typing import BinaryIO, List, Tuple, Union


COLOR_TYPE_V = 0
COLOR_TYPE_VA = 4
COLOR_TYPE_RGB = 2
COLOR_TYPE_RGBA = 6
PixelV = Tuple[int]
PixelVA = Tuple[int, int]
PixelRGB = Tuple[int, int, int]
PixelRGBA = Tuple[int, int, int, int]
ImageV = Tuple[int, List[List[PixelV]]]
ImageVA = Tuple[int, List[List[PixelVA]]]
ImageRGB = Tuple[int, List[List[PixelRGB]]]
ImageRGBA = Tuple[int, List[List[PixelRGBA]]]


def generate_checkerboard_pattern(width: int, height: int, color_type = COLOR_TYPE_RGB) -> Union[ImageV, ImageVA, ImageRGB, ImageRGBA]:
    black_pixel, white_pixel = {
        COLOR_TYPE_V: ((0,), (255,)),
        COLOR_TYPE_VA: ((0, 255,), (255, 255,)),
        COLOR_TYPE_RGB: ((0, 0, 0,), (255, 255, 255,)),
        COLOR_TYPE_RGBA: ((0, 0, 0, 255,), (255, 255, 255, 255,)),
    }[color_type]
    pixels = []
    for i in range(height):
        # Generate a single row of white/black pixels
        row = []
        for j in range(width):
            if (i + j) % 2 == 0:
                row.append(white_pixel)
            else:
                row.append(black_pixel)
        pixels.append(row)
    return (color_type, pixels)


def get_checksum(chunk_type: bytes, data: bytes) -> int:
    checksum = zlib.crc32(chunk_type)
    checksum = zlib.crc32(data, checksum)
    return checksum


def chunk(out: BinaryIO, chunk_type: bytes, data: bytes) -> None:
    out.write(struct.pack('>I', len(data)))
    out.write(chunk_type)
    out.write(data)
    checksum = get_checksum(chunk_type, data)
    out.write(struct.pack('>I', checksum))


def make_ihdr(width: int, height: int, bit_depth: int, color_type: int) -> bytes:
    return struct.pack('>2I5B', width, height, bit_depth, color_type, 0, 0, 0)


def encode_data(img: Union[ImageV, ImageVA, ImageRGB, ImageRGBA]) -> List[int]:
    ret = []

    for row in img[1]:
        ret.append(0) # Filter Method Type: None

        color_values = [
            color_value
            for pixel in row
            for color_value in pixel
        ]
        ret.extend(color_values)

    return ret


def compress_data(data: List[int]) -> bytes:
    data_bytes = bytearray(data)
    return zlib.compress(data_bytes, level=9)


def make_idat(img: Union[ImageV, ImageVA, ImageRGB, ImageRGBA]) -> bytes:
    encoded_data = encode_data(img)
    compressed_data = compress_data(encoded_data)
    return compressed_data


def dump_png(out: BinaryIO, img: Union[ImageV, ImageVA, ImageRGB, ImageRGBA]) -> None:
    HEADER = b'\x89PNG\x0d\x0a\x1a\x0a'
    out.write(HEADER)  # start by writing the header

    assert len(img[1]) > 0  # assume we were not given empty image data
    width = len(img[1][0])
    height = len(img[1])
    bit_depth = 8  # bits per pixel
    color_type = img[0]

    ihdr_data = make_ihdr(width, height, bit_depth, color_type)
    chunk(out, b'IHDR', ihdr_data)

    compressed_data = make_idat(img)
    chunk(out, b'IDAT', compressed_data)

    chunk(out, b'IEND', b'')


if __name__ == '__main__':
    width = 64
    height = 64
    for color_type, fname in [(COLOR_TYPE_V, 'out_v.png'), (COLOR_TYPE_VA, 'out_va.png'), (COLOR_TYPE_RGB, 'out_rgb.png'), (COLOR_TYPE_RGBA, 'out_rgba.png')]:
        with io.BytesIO() as membuf:
            dump_png(membuf, generate_checkerboard_pattern(width, height, color_type))
            with open(fname, 'wb') as out:
                out.write(membuf.getvalue())
            print((b'data:image/png;base64,' + base64.b64encode(membuf.getvalue())).decode('utf-8'))