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 params.add_argument( 118 '--custom-dictionary', 119 metavar='FILE', 120 type=str, 121 dest='dictfile', 122 help='Custom dictionary file.', 123 default=None) 124 # set default values using global DEFAULT_PARAMS dictionary 125 parser.set_defaults(**DEFAULT_PARAMS) 126 127 options = parser.parse_args(args=args) 128 129 if options.infile: 130 if not os.path.isfile(options.infile): 131 parser.error('file "%s" not found' % options.infile) 132 with open(options.infile, 'rb') as infile: 133 data = infile.read() 134 else: 135 if sys.stdin.isatty(): 136 # interactive console, just quit 137 parser.error('no input') 138 infile = get_binary_stdio('stdin') 139 data = infile.read() 140 141 if options.outfile: 142 if os.path.isfile(options.outfile) and not options.force: 143 parser.error('output file exists') 144 outfile = open(options.outfile, 'wb') 145 else: 146 outfile = get_binary_stdio('stdout') 147 148 if options.dictfile: 149 if not os.path.isfile(options.dictfile): 150 parser.error('file "%s" not found' % options.dictfile) 151 with open(options.dictfile, 'rb') as dictfile: 152 custom_dictionary = dictfile.read() 153 else: 154 custom_dictionary = '' 155 156 try: 157 if options.decompress: 158 data = brotli.decompress(data, dictionary=custom_dictionary) 159 else: 160 data = brotli.compress( 161 data, 162 mode=options.mode, 163 quality=options.quality, 164 lgwin=options.lgwin, 165 lgblock=options.lgblock, 166 dictionary=custom_dictionary) 167 except brotli.error as e: 168 parser.exit(1, 169 'bro: error: %s: %s' % (e, options.infile or 'sys.stdin')) 170 171 outfile.write(data) 172 outfile.close() 173 174 175if __name__ == '__main__': 176 main() 177