brotli/python/bro.py

#! /usr/bin/env python
"""bro %s -- compression/decompression utility using the Brotli algorithm."""

from __future__ import print_function
import argparse
import sys
import os
import brotli
import platform


# default values of encoder parameters
DEFAULT_PARAMS = {
    'mode': brotli.MODE_GENERIC,
    'quality': 11,
    'lgwin': 22,
    'lgblock': 0,
}


def get_binary_stdio(stream):
    """ Return the specified standard input, output or errors stream as a
    'raw' buffer object suitable for reading/writing binary data from/to it.
    """
    assert stream in ['stdin', 'stdout', 'stderr'], "invalid stream name"
    stdio = getattr(sys, stream)
    if sys.version_info[0] < 3:
        if sys.platform == 'win32':
            # set I/O stream binary flag on python2.x (Windows)
            runtime = platform.python_implementation()
            if runtime == "PyPy":
                # the msvcrt trick doesn't work in pypy, so I use fdopen
                mode = "rb" if stream == "stdin" else "wb"
                stdio = os.fdopen(stdio.fileno(), mode, 0)
            else:
                # this works with CPython -- untested on other implementations
                import msvcrt
                msvcrt.setmode(stdio.fileno(), os.O_BINARY)
        return stdio
    else:
        # get 'buffer' attribute to read/write binary data on python3.x
        if hasattr(stdio, 'buffer'):
            return stdio.buffer
        else:
            orig_stdio = getattr(sys, "__%s__" % stream)
            return orig_stdio.buffer


def main(args=None):

    parser = argparse.ArgumentParser(
        prog='bro.py',
        description="Compression/decompression utility using the Brotli algorithm.")
    parser.add_argument('--version', action='version', version=brotli.__version__)
    parser.add_argument('-i', '--input', metavar='FILE', type=str, dest='infile',
                        help='Input file', default=None)
    parser.add_argument('-o', '--output', metavar='FILE', type=str, dest='outfile',
                        help='Output file', default=None)
    parser.add_argument('-f', '--force', action='store_true',
                        help='Overwrite existing output file', default=False)
    parser.add_argument('-d', '--decompress', action='store_true',
                        help='Decompress input file', default=False)
    params = parser.add_argument_group('optional encoder parameters')
    params.add_argument('-m', '--mode', metavar="MODE", type=int, choices=[0, 1, 2],
                        help='The compression mode can be 0 for generic input, '
                        '1 for UTF-8 encoded text, or 2 for WOFF 2.0 font data. '
                        'Defaults to 0.')
    params.add_argument('-q', '--quality', metavar="QUALITY", type=int,
                        choices=list(range(0, 12)),
                        help='Controls the compression-speed vs compression-density '
                        'tradeoff. The higher the quality, the slower the '
                        'compression. Range is 0 to 11. Defaults to 11.')
    params.add_argument('--lgwin', metavar="LGWIN", type=int,
                        choices=list(range(10, 25)),
                        help='Base 2 logarithm of the sliding window size. Range is '
                        '10 to 24. Defaults to 22.')
    params.add_argument('--lgblock', metavar="LGBLOCK", type=int,
                        choices=[0] + list(range(16, 25)),
                        help='Base 2 logarithm of the maximum input block size. '
                        'Range is 16 to 24. If set to 0, the value will be set based '
                        'on the quality. Defaults to 0.')
    params.add_argument('--custom-dictionary', metavar="FILE", type=str, dest='dictfile',
                        help='Custom dictionary file.', default = None)
    # set default values using global DEFAULT_PARAMS dictionary
    parser.set_defaults(**DEFAULT_PARAMS)

    options = parser.parse_args(args=args)

    if options.infile:
        if not os.path.isfile(options.infile):
            parser.error('file "%s" not found' % options.infile)
        with open(options.infile, "rb") as infile:
            data = infile.read()
    else:
        if sys.stdin.isatty():
            # interactive console, just quit
            parser.error('no input')
        infile = get_binary_stdio('stdin')
        data = infile.read()

    if options.outfile:
        if os.path.isfile(options.outfile) and not options.force:
            parser.error('output file exists')
        outfile = open(options.outfile, "wb")
    else:
        outfile = get_binary_stdio('stdout')

    if options.dictfile:
        if not os.path.isfile(options.dictfile):
            parser.error('file "%s" not found' % options.dictfile)
        with open(options.dictfile, "rb") as dictfile:
            custom_dictionary = dictfile.read()
    else:
        custom_dictionary = ''


    try:
        if options.decompress:
            data = brotli.decompress(data, dictionary=custom_dictionary)
        else:
            data = brotli.compress(
                data, mode=options.mode, quality=options.quality,
                lgwin=options.lgwin, lgblock=options.lgblock, dictionary=custom_dictionary)
    except brotli.error as e:
        parser.exit(1,'bro: error: %s: %s' % (e, options.infile or 'sys.stdin'))

    outfile.write(data)
    outfile.close()


if __name__ == '__main__':
    main()