<commit_before>from collections import namedtuple
import zlib

class InvalidPNG(Exception):
    pass

class UnderConstruction(Exception):
    ''' Imagine the little animated GIF
    '''
    pass

# TODO: Make dicts to hold the different pixel types 
# and the different filter types

def split_into_chunks(file_bytes, chunks=[]):
    ''' chunks should look like this:
        {'length': int, 'type': str, 'data': bytes, 'crc': int}
    '''
    if not file_bytes:
        return chunks

    chunk_data_length = int.from_bytes(file_bytes[0:4], 'big')
    data_end = 8 + chunk_data_length # 8 bytes for the length and the type
    chunk_end = data_end + 4 # The 4 bytes represent the crc
    chunk_type = file_bytes[4:8].decode('utf-8')
    chunk_data = file_bytes[8:data_end]
    chunck_crc = file_bytes[data_end:chunk_end]

    chunk = {
        'length': chunk_data_length,
        'type': chunk_type,
        'data': chunk_data,
        'crc': chunck_crc
    }

    chunks.append(chunk)

    return split_into_chunks(file_bytes[chunk_end:], chunks)

def parse_ihdr_data(ihdr_chunk):
    ihdr_data = ihdr_chunk.get('data')
    Header = namedtuple(
        'Header',
        ['width', 'height', 'bit_depth', 'color_type','compression_type',
        'filter_type', 'interlace_type'])    

    return Header(
        width=int.from_bytes(ihdr_data[0:4], 'big'),
        height=int.from_bytes(ihdr_data[4:8], 'big'),
        bit_depth=ihdr_data[8],
        color_type=ihdr_data[9],
        compression_type=ihdr_data[10],
        filter_type = ihdr_data[11],
        interlace_type = ihdr_data[12]
    )

def parse_chunks(chunks):
    try:
        image_header = parse_ihdr_data([
            chunk for chunk in chunks 
            if chunk.get('type') == 'IHDR'
        ][0])
    except IndexError as err:
        raise InvalidPNG('Missing IHDR chunk')

    try:
        image_data = zlib.decompress(
            b''.join([
                idat.get('data')
                for idat in [
                    chunk for chunk in chunks 
                    if chunk.get('type') == 'IDAT']]))
    except IndexError as err:
        raise InvalidPNG('Invalid IDAT chunks')

    return image_header, image_data

def split_scanlines(width, height, pixel_size, data):
    scanline_length = width * pixel_size + 1

    scanlines = [
        data[(scanline_length * i):(scanline_length * (i + 1))]
        for i in range(height)
    ]

    return scanlines

def create_pixel_scanlines(header, data):
    if header.color_type == 6:
        Pixel = namedtuple('Pixel', ['red', 'green', 'blue', 'alpha'])
    else:
        # TODO: implement other color types
        raise UnderConstruction('Sorry dude.')

    scanlines = split_scanlines(
        header.width, header.height, len(Pixel._fields), data)

    return [create_pixels(Pixel, scanline) for scanline in scanlines]

if __name__ == '__main__':
    with open('ok.png', 'rb') as file:
        image = file.read()

    valid_png_header = b'\x89PNG\r\n\x1a\n'

    if image[0:8] != valid_png_header:
        raise InvalidPNG('not a valid header')

    image_header, image_data = parse_chunks(split_into_chunks(image[8:]))

    pixels = create_pixel_scanlines(image_header, image_data)
<commit_msg>big old mess. but retiring<commit_after>from collections import namedtuple
import zlib

class InvalidPNG(Exception):
    pass

# TODO: Make dict to hold the different filter types

PIXELS = {
    0: namedtuple('Pixel', ['gray']),
    2: namedtuple('Pixel', ['red', 'green', 'blue']),
    # 3: A Palette index,
    4: namedtuple('Pixel', ['gray', 'alpha']),
    6: namedtuple('Pixel', ['red', 'green', 'blue', 'alpha'])
}

# http://www.libpng.org/pub/png/spec/1.2/PNG-Filters.html

def filter_algorithm(type, current_byte, index_in_scanline, scanline, prior_scanline=[]):
    if type == 0:
        return current_byte
    elif type == 1:
        ''' Sub[x] + Raw(x-bpp)
            Sub(x) is the current_byte
            Raw(x-bpp) is the previous byte (of that color)
        '''
        sub_x = current_byte
        raw_x_bpp = 0 if (index_in_scanline - bytes_per_pixel) <= bytes_per_pixel else index_in_scanline - bytes_per_pixel
        # print(int(current_byte + scanline[raw_x_bpp]))
        return int(current_byte + scanline[raw_x_bpp]) % 256

# def none(current_byte, index_in_scanline, scanline, prior_scanline=[]):
#     return current_byte

# def sub(current_byte, index_in_scanline, scanline, prior_scanline=[]):
#     ''' Sub(x) + Raw(x-bpp)
#     '''
#     if index_in_scanline == 0:

#     return 
# def up(current_byte, prior_byte, byte_above, byte_prior_to_above, bytes_per_pixel):
#     pass
# def average(current_byte, prior_byte, byte_above, byte_prior_to_above, bytes_per_pixel):
#     pass
# def paeth(current_byte, prior_byte, byte_above, byte_prior_to_above, bytes_per_pixel):
#     pass

# FILTER_ALGORITHMS = {
#     0: lambda *args: None,
#     1: sub,
#     2: up,
#     3: average,
#     4: paeth
# }

def split_into_chunks(file_bytes, chunks=[]):
    ''' chunks should look like this:
        {'length': int, 'type': str, 'data': bytes, 'crc': int}
    '''
    if not file_bytes:
        return chunks

    chunk_data_length = int.from_bytes(file_bytes[0:4], 'big')
    data_end = 8 + chunk_data_length # 8 bytes for the length and the type
    chunk_end = data_end + 4 # The 4 bytes represent the crc
    chunk_type = file_bytes[4:8].decode('utf-8')
    chunk_data = file_bytes[8:data_end]
    chunck_crc = file_bytes[data_end:chunk_end]

    chunk = {
        'length': chunk_data_length,
        'type': chunk_type,
        'data': chunk_data,
        'crc': chunck_crc
    }

    chunks.append(chunk)

    return split_into_chunks(file_bytes[chunk_end:], chunks)

def parse_ihdr_data(ihdr_chunk):
    ihdr_data = ihdr_chunk.get('data')
    Header = namedtuple(
        'Header',
        ['width', 'height', 'bit_depth', 'color_type','compression_type',
        'filter_type', 'interlace_type'])    

    return Header(
        width=int.from_bytes(ihdr_data[0:4], 'big'),
        height=int.from_bytes(ihdr_data[4:8], 'big'),
        bit_depth=ihdr_data[8],
        color_type=ihdr_data[9],
        compression_type=ihdr_data[10],
        filter_type = ihdr_data[11],
        interlace_type = ihdr_data[12]
    )

def parse_chunks(chunks):
    try:
        image_header = parse_ihdr_data([
            chunk for chunk in chunks 
            if chunk.get('type') == 'IHDR'
        ][0])
    except IndexError as err:
        raise InvalidPNG('Missing IHDR chunk')

    try:
        image_data = zlib.decompress(
            b''.join([
                idat.get('data')
                for idat in [
                    chunk for chunk in chunks 
                    if chunk.get('type') == 'IDAT']]))
    except IndexError as err:
        raise InvalidPNG('Invalid IDAT chunks')

    return image_header, image_data

def split_raw_scanlines(width, height, bytes_per_pixel, data):
    scanline_length = width * bytes_per_pixel + 1

    scanlines = [
        data[(scanline_length * i):(scanline_length * (i + 1))]
        for i in range(height)
    ]

    return scanlines

def create_pixels(Pixel, scanline, bytes_per_pixel, prior_scanline=[]):
    filter_type_byte = scanline[0]
    bytes_only = scanline[1:]

    return [
        filter_algorithm(
            filter_type_byte,
            byte,
            i, 
            bytes_only
        )
        for i, byte in enumerate(bytes_only)
    ]


    # algo = FILTER_ALGORITHMS[filter_type_byte]

    # [algo(byte, i, scanline) for i, byte in enumerate(scanline)]

    # data = scanline[1:]
    # pixel_range = range(len(scanline[1:]))[::len(Pixel._fields)]

    # return [
    #     Pixel(*data[pixel_index:pixel_range[i+1]])
    #     for i, pixel_index in enumerate(pixel_range[:-1])]

if __name__ == '__main__':
    with open('test.png', 'rb') as file:
        image = file.read()

    valid_png_header = b'\x89PNG\r\n\x1a\n'

    if image[0:8] != valid_png_header:
        raise InvalidPNG('not a valid header')

    image_header, image_data = parse_chunks(split_into_chunks(image[8:]))

    try:
        Pixel = PIXELS[image_header.color_type]
    except KeyError as err:
        raise UnemplementedError('I haven\'t done that yet.')
    else:
        bytes_per_pixel = int(
            len(Pixel._fields) * (image_header.bit_depth / 8))

    scanlines = split_raw_scanlines(
        image_header.width,
        image_header.height,
        bytes_per_pixel, 
        image_data
    )

    pixels = create_pixels(Pixel, scanlines[0], bytes_per_pixel)

    # pixel_lines = create_scanlines(image_header, image_data)

{'fraction_non_alphanumeric': 0.1205000578770691, 'fraction_numerical': 0.010417872438939692, 'mean_word_length': 2.819628647214854, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '429773', 'n_tokens_mistral': 2980, 'n_tokens_neox': 2881, 'n_words': 646}