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