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