imgdoc

#! /usr/bin/env python3
from argparse import ArgumentParser
from PIL import Image
import sys
import zlib
import os


def yesno(question):
    answered = False
    answer = False
    while not answered:
        ans = input(f'{question} (y/N): ').lower()
        if ans in ['y', 'yes', 'n', 'no']:
            answered = True
        if ans in ['y', 'yes']:
            answer = True
    return answer


def main():
    parser = ArgumentParser(
        description='ImgDoc, convert a DocImg PNG back into its original format.')
    parser.add_argument('input', type=str,
                        help='The input document to convert back.')
    parser.add_argument('output', type=str,
                        help='The output document to save data to.')
    args = parser.parse_args()
    inputf = args.input
    outf = args.output
    split_outf = os.path.split(outf)
    if not os.path.isfile(inputf):
        parser.error(f'Couldn\'t find file: {inputf}')
    if split_outf[0].strip() and not os.path.isdir(split_outf[0]):
        parser.error(f'Directory doesn\'t exist: {split_outf[0]}')
    print('Reading input file...')
    image = Image.open(inputf)
    print('Done reading.')
    print('Getting colors...')
    color_pairs = list(image.getdata())
    print('Done getting colors.')
    print('Building back data...')
    data = bytearray()
    for pair in color_pairs:
        data.append(pair[0])
    print('Done building data.')
    print('Decompressing data...')
    dc_data = zlib.decompress(data)
    print('Done decompressing data.')
    print('Saving data...')
    if os.path.isfile(outf):
        if yesno(f'{outf} already exists. Replace?'):
            print('Replacing output target.')
        else:
            sys.exit('Aborting.')
    with open(outf, 'wb') as outfile:
        print('Writing data...')
        outfile.write(dc_data)
        print('Done writing data.')
    print('Done saving data.')


if __name__ == '__main__':
    main()