import zlib
import struct
import array

def write_png(filename, width, height, rgb_func):
    
    def output_chunk(out, chunk_type, data):
        out.write(struct.pack("!I", len(data)))
        out.write(chunk_type)
        out.write(data)
        checksum = zlib.crc32(data, zlib.crc32(chunk_type))
        out.write(struct.pack("!I", checksum))
    
    def get_data(width, height, rgb_func):
        compressor = zlib.compressobj()
        data = array.array("B")
        for y in range(height):
            data.append(0)
            for x in range(width):
                data.extend(rgb_func(x, y))
        compressed = compressor.compress(data.tostring())
        flushed = compressor.flush()
        return compressed + flushed
    
    out = open(filename, "w")
    out.write(struct.pack("8B", 137, 80, 78, 71, 13, 10, 26, 10))
    output_chunk(out, "IHDR", struct.pack("!2I5B", width, height, 8, 2, 0, 0, 0))
    output_chunk(out, "IDAT", get_data(width, height, rgb_func))
    output_chunk(out, "IEND", "")
    out.close()