From ab303586fcd2e55e6459e204ca0988846d1d2afd Mon Sep 17 00:00:00 2001 From: drewcassidy Date: Sun, 11 Apr 2021 20:49:31 -0700 Subject: [PATCH] command line tool --- pillow_mbm/__init__.py | 40 ++++++++++++++++++- pillow_mbm/__main__.py | 90 +++++++++++++++++++++++++++++++++++++++++- setup.py | 30 ++++++-------- 3 files changed, 140 insertions(+), 20 deletions(-) diff --git a/pillow_mbm/__init__.py b/pillow_mbm/__init__.py index a12e1ba..b65c368 100644 --- a/pillow_mbm/__init__.py +++ b/pillow_mbm/__init__.py @@ -1 +1,39 @@ -from version import __version__ +from pillow_mbm.version import __version__ + +import struct +from PIL import Image, ImageFile + +MAGIC = b'\x03KSP' + + +def _accept(prefix: bytes): + """Check if a file is a MBM file""" + return prefix[:4] == MAGIC + + +class MBMImageFile(ImageFile.ImageFile): + format = 'MBM' + format_description = 'Kerbal Space Program MBM image' + + def _open(self): + """Open an MBM file""" + magic = self.fp.read(4) + if magic != MAGIC: + raise SyntaxError('not a MBM file') + + width, height, bits = struct.unpack('<2I4xI', self.fp.read(16)) + + self._size = (width, height) + + if bits == 24: + self.mode = 'RGB' + elif bits == 32: + self.mode = 'RGBA' + else: + raise SyntaxError('unknown number of bits') + + self.tile = [('raw', (0, 0, width, height), 20, (self.mode, 0, 1))] + + +Image.register_open(MBMImageFile.format, MBMImageFile, _accept) +Image.register_extensions(MBMImageFile.format, ['.mbm']) diff --git a/pillow_mbm/__main__.py b/pillow_mbm/__main__.py index 100915d..2fbc747 100644 --- a/pillow_mbm/__main__.py +++ b/pillow_mbm/__main__.py @@ -1,2 +1,88 @@ -def main(): - print('theres nothing here yet!') +import pathlib +import os +import click +from typing import List +from PIL import Image + + +def get_decoded_extensions(feature: str = 'open') -> List[str]: + """Gets a list of extensions for Pillow formats supporting a supplied feature""" + Image.init() + extensions = Image.EXTENSION + formats = getattr(Image, feature.upper()).keys() + + return [ext for ext, fmt in extensions.items() if fmt in formats] + + +decoded_extensions = get_decoded_extensions('save') + + +# noinspection PyUnusedLocal +def validate_decoded_extension(ctx, param, value) -> str: + """Check if an extension for a decoded image is valid""" + if value[0] != '.': + value = '.' + value + + if value not in decoded_extensions: + raise click.BadParameter(f'Invalid extension for decoded file. Valid extensions are:\n{decoded_extensions}') + + return value + + +def path_pairs(inputs, output, suffix, extension): + if len(inputs) < 1: + raise click.BadArgumentUsage('No input files were provided.') + + inpaths = [pathlib.Path(i) for i in inputs] + + if not output: + # decode in place + return [(inpath, inpath.with_name(inpath.stem + suffix + extension)) for inpath in inpaths] + + else: + outpath = pathlib.Path(output) + if outpath.is_file(): + # decode to a file + if len(inputs) > 1: + raise click.BadOptionUsage('output', 'Output is a single file, but multiple input files were provided.') + if outpath.suffix not in decoded_extensions: + raise click.BadOptionUsage('output', f'File has incorrect extension for decoded file. Valid extensions are:\n{decoded_extensions}') + + return [(inpath, outpath) for inpath in inpaths] + else: + # decode to directory + return [(inpath, outpath / (inpath.stem + suffix + extension)) for inpath in inpaths] + + +@click.command() +@click.option('-f/-F', '--flip/--no-flip', default=True, show_default=False, help="Vertically flip image after converting.") +@click.option('-r', '--remove', is_flag=True, help="Remove input images after converting.") +@click.option('-s', '--suffix', type=str, default='', help="Suffix to append to output file(s). Ignored if output is a single file.") +@click.option('-x', '--extension', + callback=validate_decoded_extension, + type=str, default='.png', show_default=True, + help="Extension to use for output. Ignored if output is a single file. Output filetype is deduced from this") +@click.option('-o', '--output', + type=click.Path(writable=True), default=None, + help="Output file or directory. If outputting to a file, input filenames must be only a single item. By default, files are decoded in place.") +@click.argument('filenames', nargs=-1, type=click.Path(exists=True, readable=True, dir_okay=False)) +def decode(flip, remove, suffix, extension, output, filenames): + """Decode Kerbal Space Program MBM files""" + + pairs = path_pairs(filenames, output, suffix, extension) + + with click.progressbar(pairs, show_eta=False, show_pos=True, item_show_func=lambda x: str(x[0]) if x else '') as bar: + for inpath, outpath in bar: + image = Image.open(inpath) + + if flip: + image = image.transpose(Image.FLIP_TOP_BOTTOM) + + image.save(outpath) + + if remove: + os.remove(inpath) + + +if __name__ == '__main__': + decode() diff --git a/setup.py b/setup.py index 6dd1f8c..2226fdb 100644 --- a/setup.py +++ b/setup.py @@ -1,32 +1,28 @@ +import os.path from setuptools import setup +project_path = os.path.dirname(os.path.realpath(__file__)) -def version(): - with open('pillow_mbm/version.py') as f: - exec(f.read()) - return __version__ - - -def readme(): - with open('README.md') as f: - return f.read() +with open(os.path.join(project_path, 'pillow_mbm', 'version.py')) as f: + exec(f.read()) +with open(os.path.join(project_path, 'README.md')) as f: + readme = f.read() setup( name='pillow-mbm', description="A pillow plugin that adds support for KSP's MBM textures", - version=version(), - long_description=readme(), + version=__version__, + long_description=readme, long_description_content_type='text/markdown', + python_requires=">=3.7", install_requires=['Pillow'], - extras_require={ - 'cli': ['click'] + extras_require={'CLI': ['click']}, + entry_points={ + 'console_scripts': ['convert-mbm = pillow_mbm.__main__:decode [CLI]'] }, package_dir={'': '.'}, - entry_points=''' - [console_scripts] - convert-mbm=pillow_mbm.__main__:main [cli] - ''', + packages=['pillow_mbm'], classifiers=[ 'Development Status :: 1 - Planning', 'Intended Audience :: Developers',