1#! /usr/bin/env python
2"""Compression/decompression utility using the Brotli algorithm."""
3
4from __future__ import print_function
5import argparse
6import sys
7import os
8import platform
9
10import brotli
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=os.path.basename(__file__), description=__doc__)
53    parser.add_argument(
54        '--version', action='version', version=brotli.__version__)
55    parser.add_argument(
56        '-i',
57        '--input',
58        metavar='FILE',
59        type=str,
60        dest='infile',
61        help='Input file',
62        default=None)
63    parser.add_argument(
64        '-o',
65        '--output',
66        metavar='FILE',
67        type=str,
68        dest='outfile',
69        help='Output file',
70        default=None)
71    parser.add_argument(
72        '-f',
73        '--force',
74        action='store_true',
75        help='Overwrite existing output file',
76        default=False)
77    parser.add_argument(
78        '-d',
79        '--decompress',
80        action='store_true',
81        help='Decompress input file',
82        default=False)
83    params = parser.add_argument_group('optional encoder parameters')
84    params.add_argument(
85        '-m',
86        '--mode',
87        metavar='MODE',
88        type=int,
89        choices=[0, 1, 2],
90        help='The compression mode can be 0 for generic input, '
91        '1 for UTF-8 encoded text, or 2 for WOFF 2.0 font data. '
92        'Defaults to 0.')
93    params.add_argument(
94        '-q',
95        '--quality',
96        metavar='QUALITY',
97        type=int,
98        choices=list(range(0, 12)),
99        help='Controls the compression-speed vs compression-density '
100        'tradeoff. The higher the quality, the slower the '
101        'compression. Range is 0 to 11. Defaults to 11.')
102    params.add_argument(
103        '--lgwin',
104        metavar='LGWIN',
105        type=int,
106        choices=list(range(10, 25)),
107        help='Base 2 logarithm of the sliding window size. Range is '
108        '10 to 24. Defaults to 22.')
109    params.add_argument(
110        '--lgblock',
111        metavar='LGBLOCK',
112        type=int,
113        choices=[0] + list(range(16, 25)),
114        help='Base 2 logarithm of the maximum input block size. '
115        'Range is 16 to 24. If set to 0, the value will be set based '
116        'on the quality. Defaults to 0.')
117    # set default values using global DEFAULT_PARAMS dictionary
118    parser.set_defaults(**DEFAULT_PARAMS)
119
120    options = parser.parse_args(args=args)
121
122    if options.infile:
123        if not os.path.isfile(options.infile):
124            parser.error('file "%s" not found' % options.infile)
125        with open(options.infile, 'rb') as infile:
126            data = infile.read()
127    else:
128        if sys.stdin.isatty():
129            # interactive console, just quit
130            parser.error('no input')
131        infile = get_binary_stdio('stdin')
132        data = infile.read()
133
134    if options.outfile:
135        if os.path.isfile(options.outfile) and not options.force:
136            parser.error('output file exists')
137        outfile = open(options.outfile, 'wb')
138    else:
139        outfile = get_binary_stdio('stdout')
140
141    try:
142        if options.decompress:
143            data = brotli.decompress(data)
144        else:
145            data = brotli.compress(
146                data,
147                mode=options.mode,
148                quality=options.quality,
149                lgwin=options.lgwin,
150                lgblock=options.lgblock)
151    except brotli.error as e:
152        parser.exit(1,
153                    'bro: error: %s: %s' % (e, options.infile or 'sys.stdin'))
154
155    outfile.write(data)
156    outfile.close()
157
158
159if __name__ == '__main__':
160    main()
161