Home | History | Annotate | Download | only in python 
      1 #! /usr/bin/env python
      2 """Compression/decompression utility using the Brotli algorithm."""
      3 
      4 from __future__ import print_function
      5 import argparse
      6 import sys
      7 import os
      8 import platform
      9 
     10 import brotli
     11 
     12 # default values of encoder parameters
     13 DEFAULT_PARAMS = {
     14     'mode': brotli.MODE_GENERIC,
     15     'quality': 11,
     16     'lgwin': 22,
     17     'lgblock': 0,
     18 }
     19 
     20 
     21 def 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 
     49 def 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 
    159 if __name__ == '__main__':
    160     main()
    161