bro.py revision 291d21cc2ffe042fc3134d935e4f810184e82138 
1#! /usr/bin/env python
2"""bro %s -- compression/decompression utility using the Brotli algorithm."""
3
4from __future__ import print_function
5import argparse
6import sys
7import os
8import brotli
9import platform
10
11
12# default values of encoder parameters
13DEFAULT_PARAMS = {
14    'mode': brotli.MODE_GENERIC,
15    'quality': 11,
16    'lgwin': 22,
17    'lgblock': 0,
18}
19
20
21def get_binary_stdio(stream):
22    """ Return the specified standard input, output or errors stream as a
23    'raw' buffer object suitable for reading/writing binary data from/to it.
24    """
25    assert stream in ['stdin', 'stdout', 'stderr'], "invalid stream name"
26    stdio = getattr(sys, stream)
27    if sys.version_info[0] < 3:
28        if sys.platform == 'win32':
29            # set I/O stream binary flag on python2.x (Windows)
30            runtime = platform.python_implementation()
31            if runtime == "PyPy":
32                # the msvcrt trick doesn't work in pypy, so I use fdopen
33                mode = "rb" if stream == "stdin" else "wb"
34                stdio = os.fdopen(stdio.fileno(), mode, 0)
35            else:
36                # this works with CPython -- untested on other implementations
37                import msvcrt
38                msvcrt.setmode(stdio.fileno(), os.O_BINARY)
39        return stdio
40    else:
41        # get 'buffer' attribute to read/write binary data on python3.x
42        if hasattr(stdio, 'buffer'):
43            return stdio.buffer
44        else:
45            orig_stdio = getattr(sys, "__%s__" % stream)
46            return orig_stdio.buffer
47
48
49def main(args=None):
50
51    parser = argparse.ArgumentParser(
52        prog='bro.py',
53        description="Compression/decompression utility using the Brotli algorithm.")
54    parser.add_argument('--version', action='version', version=brotli.__version__)
55    parser.add_argument('-i', '--input', metavar='FILE', type=str, dest='infile',
56                        help='Input file', default=None)
57    parser.add_argument('-o', '--output', metavar='FILE', type=str, dest='outfile',
58                        help='Output file', default=None)
59    parser.add_argument('-f', '--force', action='store_true',
60                        help='Overwrite existing output file', default=False)
61    parser.add_argument('-d', '--decompress', action='store_true',
62                        help='Decompress input file', default=False)
63    params = parser.add_argument_group('optional encoder parameters')
64    params.add_argument('-m', '--mode', metavar="MODE", type=int, choices=[0, 1, 2],
65                        help='The compression mode can be 0 for generic input, '
66                        '1 for UTF-8 encoded text, or 2 for WOFF 2.0 font data. '
67                        'Defaults to 0.')
68    params.add_argument('-q', '--quality', metavar="QUALITY", type=int,
69                        choices=list(range(0, 12)),
70                        help='Controls the compression-speed vs compression-density '
71                        'tradeoff. The higher the quality, the slower the '
72                        'compression. Range is 0 to 11. Defaults to 11.')
73    params.add_argument('--lgwin', metavar="LGWIN", type=int,
74                        choices=list(range(10, 25)),
75                        help='Base 2 logarithm of the sliding window size. Range is '
76                        '10 to 24. Defaults to 22.')
77    params.add_argument('--lgblock', metavar="LGBLOCK", type=int,
78                        choices=[0] + list(range(16, 25)),
79                        help='Base 2 logarithm of the maximum input block size. '
80                        'Range is 16 to 24. If set to 0, the value will be set based '
81                        'on the quality. Defaults to 0.')
82    # set default values using global DEFAULT_PARAMS dictionary
83    parser.set_defaults(**DEFAULT_PARAMS)
84
85    options = parser.parse_args(args=args)
86
87    if options.infile:
88        if not os.path.isfile(options.infile):
89            parser.error('file "%s" not found' % options.infile)
90        with open(options.infile, "rb") as infile:
91            data = infile.read()
92    else:
93        if sys.stdin.isatty():
94            # interactive console, just quit
95            parser.error('no input')
96        infile = get_binary_stdio('stdin')
97        data = infile.read()
98
99    if options.outfile:
100        if os.path.isfile(options.outfile) and not options.force:
101            parser.error('output file exists')
102        outfile = open(options.outfile, "wb")
103    else:
104        outfile = get_binary_stdio('stdout')
105
106    try:
107        if options.decompress:
108            data = brotli.decompress(data)
109        else:
110            data = brotli.compress(
111                data, mode=options.mode, quality=options.quality,
112                lgwin=options.lgwin, lgblock=options.lgblock)
113    except brotli.error as e:
114        parser.exit(1,'bro: error: %s: %s' % (e, options.infile or 'sys.stdin'))
115
116    outfile.write(data)
117    outfile.close()
118
119
120if __name__ == '__main__':
121    main()
122