toys/map/scripts/mappkg.py

#!/usr/bin/env python3
# -*- coding: utf-8; mode: python; tab-width: 4; indent-tabs-mode: nil -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8

import sys
import os
import json
import shutil
import argparse
import textwrap
from zipfile import ZipFile, is_zipfile


__VERSION__ = '.'.join(map(str, [0, 1, 1]))


def main():
    parser = argparse.ArgumentParser(description='Create, check and uncompress the map package',
        usage='use "python %(prog)s --help for more information',
        formatter_class=argparse.RawTextHelpFormatter,
        epilog='Typical usage\n'
            "\tpython %(prog)s -o map.zip -d <dir>\n"
            "\tpython %(prog)s -x map.zip -d <dir>\n"
    	    "\tpython %(prog)s -c map.zip")
    parser.add_argument('-v', '--version', action='version', version='{}'.format(__VERSION__))
    parser.add_argument('-c', '--check',
        help=textwrap.dedent('''Check the map package according to the prefab rules.'''))
    parser.add_argument('-t', '--text', default=None,
        help=textwrap.dedent('''Provide the description information of this package, just use the file name (without extension), it is a patch for Jenkins Pipeline Script.'''))
    parser.add_argument('-o', '--output',
        help=textwrap.dedent('''Create a compressed map package.'''))
    parser.add_argument('-d', '--directory', default='.',
        help=textwrap.dedent('''Specify the directory should be operated.'''))
    parser.add_argument('-x', '--extract',
        help=textwrap.dedent('''Decompress the map package.'''))

    args = sys.argv[1:]
    if len(args) == 0:
        parser.print_help()
        return 0
    else:
        parsed_args = parser.parse_args()

    if parsed_args.check:
        file_path = os.path.abspath(parsed_args.check)
        if os.path.isfile(file_path) and is_zipfile(file_path):
            # Do checking
            dst_dir = os.path.abspath(parsed_args.directory)
            try:
                mappkg = ZipFile(file_path, 'r')
                if not os.path.exists(dst_dir):
                    os.makedirs(dst_dir)
                # Do checking
                if parsed_args.text:
                    file_base_name = parsed_args.text
                else:
                    file_base_name = os.path.splitext(os.path.basename(file_path))[0]
                image_base_name = '{}.png'.format(file_base_name)
                # The top level directory is file_base_name
                base_exists, image_exists, description_exists = False, False, False
                for item in mappkg.namelist():
                    fpath = item.split(os.sep)[0]
                    if fpath == image_base_name:
                        image_exists = True
                    elif fpath == 'description':
                        description_exists = True
                    elif fpath == file_base_name:
                        base_exists = True
                    else:
                        print("WARNING: The {} is not the vaild content".format(item))
                if not base_exists:
                    raise Exception("Required content directory <{}> is missing!".format(file_base_name))
                elif not image_exists:
                    raise Exception("Required figure image <{}> is missing".format(image_base_name))
                elif not description_exists:
                    print("WARNING: Optional long description is missing")
                mappkg.close()
            except:
                raise Exception("Failed to read the map package.")
        else:
            raise Exception("{} is not a zip file".format(file_path))
    elif parsed_args.output:
        # Create a zip file from the specified directory
        file_path = os.path.abspath(parsed_args.output)
        dst_dir = os.path.abspath(parsed_args.directory)
        os.chdir(dst_dir)
        try:
            mappkg = ZipFile(file_path, 'w')
            for root, dirs, files in os.walk('.'):
                for f in files:
                    mappkg.write(os.path.join(root, f))
        except:
            raise Exception("Failed to create map package.")
    elif parsed_args.extract:
        # Extract a map package into specified directory
        file_path = os.path.abspath(parsed_args.extract)
        if parsed_args.text:
            file_base_name = parsed_args.text
        else:
            file_base_name = os.path.splitext(os.path.basename(file_path))[0]
        image_base_name = '{}.png'.format(file_base_name)
        catalog, description, version = file_base_name.split('-')
        dst_dir = os.path.abspath(parsed_args.directory)
        try:
            mappkg = ZipFile(file_path, 'r')
            if not os.path.exists(dst_dir):
                os.makedirs(dst_dir)
            # Create convention folders
            # NOTE: Here we assume that you have check the package before this action
            map_dir = os.path.join(dst_dir, catalog, version)
            if not os.path.exists(map_dir):
                os.makedirs(map_dir)
            # Extract specify proto files into <catalog>/<version>, 
            # and extract image, description into destination folder.
            for item in mappkg.namelist():
                items = item.split('/')
                if len(items) == 2 and items[0] == file_base_name and items[1].endswith('.proto'):
                    # copy file (taken from zipfile's extract)
                    src = mappkg.open(item)
                    tgt = open(os.path.join(map_dir, items[1]), "wb")
                    with src, tgt:
                        shutil.copyfileobj(src, tgt)
                elif items[0] == image_base_name or items[0] == 'description':
                    mappkg.extract(item, dst_dir)
            # If the `description` file is not existent,
            # create one with the short description string
            if not os.path.isfile(os.path.join(dst_dir, 'description')):
                json_str = {'zhName': catalog,
                    'enName': catalog,
                    'description': description}
                open(os.path.join(dst_dir, 'description'), 'w+').write(json.dumps(json_str, indent=4))
            mappkg.close()
        except:
            raise Exception("Failed to extract map package.")

    return 0


if __name__ == '__main__':
    sys.exit(main())