bro.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. #! /usr/bin/env python
  2. """Compression/decompression utility using the Brotli algorithm."""
  3. # Note: Python2 has been deprecated long ago, but some projects out in
  4. # the wide world may still use it nevertheless. This should not
  5. # deprive them from being able to run Brotli.
  6. from __future__ import print_function
  7. import argparse
  8. import os
  9. import platform
  10. import sys
  11. import brotli
  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. def get_binary_stdio(stream):
  20. """Return the specified stdin/stdout/stderr stream.
  21. If the stdio stream requested (i.e. sys.(stdin|stdout|stderr))
  22. has been replaced with a stream object that does not have a `.buffer`
  23. attribute, this will return the original stdio stream's buffer, i.e.
  24. `sys.__(stdin|stdout|stderr)__.buffer`.
  25. Args:
  26. stream: One of 'stdin', 'stdout', 'stderr'.
  27. Returns:
  28. The stream, as a 'raw' buffer object (i.e. io.BufferedIOBase subclass
  29. instance such as io.Bufferedreader/io.BufferedWriter), suitable for
  30. reading/writing binary data from/to it.
  31. """
  32. if stream == 'stdin': stdio = sys.stdin
  33. elif stream == 'stdout': stdio = sys.stdout
  34. elif stream == 'stderr': stdio = sys.stderr
  35. else:
  36. raise ValueError('invalid stream name: %s' % (stream,))
  37. if sys.version_info[0] < 3:
  38. if sys.platform == 'win32':
  39. # set I/O stream binary flag on python2.x (Windows)
  40. runtime = platform.python_implementation()
  41. if runtime == 'PyPy':
  42. # the msvcrt trick doesn't work in pypy, so use fdopen().
  43. mode = 'rb' if stream == 'stdin' else 'wb'
  44. stdio = os.fdopen(stdio.fileno(), mode, 0)
  45. else:
  46. # this works with CPython -- untested on other implementations
  47. import msvcrt
  48. msvcrt.setmode(stdio.fileno(), os.O_BINARY)
  49. return stdio
  50. else:
  51. try:
  52. return stdio.buffer
  53. except AttributeError:
  54. # The Python reference explains
  55. # (-> https://docs.python.org/3/library/sys.html#sys.stdin)
  56. # that the `.buffer` attribute might not exist, since
  57. # the standard streams might have been replaced by something else
  58. # (such as an `io.StringIO()` - perhaps via
  59. # `contextlib.redirect_stdout()`).
  60. # We fall back to the original stdio in these cases.
  61. if stream == 'stdin': return sys.__stdin__.buffer
  62. if stream == 'stdout': return sys.__stdout__.buffer
  63. if stream == 'stderr': return sys.__stderr__.buffer
  64. assert False, 'Impossible Situation.'
  65. def main(args=None):
  66. parser = argparse.ArgumentParser(
  67. prog=os.path.basename(__file__), description=__doc__)
  68. parser.add_argument(
  69. '--version', action='version', version=brotli.version)
  70. parser.add_argument(
  71. '-i',
  72. '--input',
  73. metavar='FILE',
  74. type=str,
  75. dest='infile',
  76. help='Input file',
  77. default=None)
  78. parser.add_argument(
  79. '-o',
  80. '--output',
  81. metavar='FILE',
  82. type=str,
  83. dest='outfile',
  84. help='Output file',
  85. default=None)
  86. parser.add_argument(
  87. '-f',
  88. '--force',
  89. action='store_true',
  90. help='Overwrite existing output file',
  91. default=False)
  92. parser.add_argument(
  93. '-d',
  94. '--decompress',
  95. action='store_true',
  96. help='Decompress input file',
  97. default=False)
  98. params = parser.add_argument_group('optional encoder parameters')
  99. params.add_argument(
  100. '-m',
  101. '--mode',
  102. metavar='MODE',
  103. type=int,
  104. choices=[0, 1, 2],
  105. help='The compression mode can be 0 for generic input, '
  106. '1 for UTF-8 encoded text, or 2 for WOFF 2.0 font data. '
  107. 'Defaults to 0.')
  108. params.add_argument(
  109. '-q',
  110. '--quality',
  111. metavar='QUALITY',
  112. type=int,
  113. choices=list(range(0, 12)),
  114. help='Controls the compression-speed vs compression-density '
  115. 'tradeoff. The higher the quality, the slower the '
  116. 'compression. Range is 0 to 11. Defaults to 11.')
  117. params.add_argument(
  118. '--lgwin',
  119. metavar='LGWIN',
  120. type=int,
  121. choices=list(range(10, 25)),
  122. help='Base 2 logarithm of the sliding window size. Range is '
  123. '10 to 24. Defaults to 22.')
  124. params.add_argument(
  125. '--lgblock',
  126. metavar='LGBLOCK',
  127. type=int,
  128. choices=[0] + list(range(16, 25)),
  129. help='Base 2 logarithm of the maximum input block size. '
  130. 'Range is 16 to 24. If set to 0, the value will be set based '
  131. 'on the quality. Defaults to 0.')
  132. # set default values using global _DEFAULT_PARAMS dictionary
  133. parser.set_defaults(**_DEFAULT_PARAMS)
  134. options = parser.parse_args(args=args)
  135. if options.infile:
  136. try:
  137. with open(options.infile, 'rb') as infile:
  138. data = infile.read()
  139. except OSError:
  140. parser.error('Could not read --infile: %s' % (infile,))
  141. else:
  142. if sys.stdin.isatty():
  143. # interactive console, just quit
  144. parser.error('No input (called from interactive terminal).')
  145. infile = get_binary_stdio('stdin')
  146. data = infile.read()
  147. if options.outfile:
  148. # Caution! If `options.outfile` is a broken symlink, will try to
  149. # redirect the write according to symlink.
  150. if os.path.exists(options.outfile) and not options.force:
  151. parser.error(('Target --outfile=%s already exists, '
  152. 'but --force was not requested.') % (outfile,))
  153. outfile = open(options.outfile, 'wb')
  154. did_open_outfile = True
  155. else:
  156. outfile = get_binary_stdio('stdout')
  157. did_open_outfile = False
  158. try:
  159. try:
  160. if options.decompress:
  161. data = brotli.decompress(data)
  162. else:
  163. data = brotli.compress(
  164. data,
  165. mode=options.mode,
  166. quality=options.quality,
  167. lgwin=options.lgwin,
  168. lgblock=options.lgblock)
  169. outfile.write(data)
  170. finally:
  171. if did_open_outfile: outfile.close()
  172. except brotli.error as e:
  173. parser.exit(1,
  174. 'bro: error: %s: %s' % (e, options.infile or '{stdin}'))
  175. if __name__ == '__main__':
  176. main()