cli.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242
  1. """
  2. Module version for monitoring CLI pipes (`... | python -m tqdm | ...`).
  3. """
  4. from .std import tqdm, TqdmTypeError, TqdmKeyError
  5. from ._version import __version__ # NOQA
  6. import sys
  7. import re
  8. import logging
  9. __all__ = ["main"]
  10. def cast(val, typ):
  11. log = logging.getLogger(__name__)
  12. log.debug((val, typ))
  13. if " or " in typ:
  14. for t in typ.split(" or "):
  15. try:
  16. return cast(val, t)
  17. except TqdmTypeError:
  18. pass
  19. raise TqdmTypeError(val + ' : ' + typ)
  20. # sys.stderr.write('\ndebug | `val:type`: `' + val + ':' + typ + '`.\n')
  21. if typ == 'bool':
  22. if (val == 'True') or (val == ''):
  23. return True
  24. elif val == 'False':
  25. return False
  26. else:
  27. raise TqdmTypeError(val + ' : ' + typ)
  28. try:
  29. return eval(typ + '("' + val + '")')
  30. except:
  31. if typ == 'chr':
  32. return chr(ord(eval('"' + val + '"')))
  33. else:
  34. raise TqdmTypeError(val + ' : ' + typ)
  35. def posix_pipe(fin, fout, delim='\n', buf_size=256,
  36. callback=lambda int: None # pragma: no cover
  37. ):
  38. """
  39. Params
  40. ------
  41. fin : file with `read(buf_size : int)` method
  42. fout : file with `write` (and optionally `flush`) methods.
  43. callback : function(int), e.g.: `tqdm.update`
  44. """
  45. fp_write = fout.write
  46. # tmp = ''
  47. if not delim:
  48. while True:
  49. tmp = fin.read(buf_size)
  50. # flush at EOF
  51. if not tmp:
  52. getattr(fout, 'flush', lambda: None)() # pragma: no cover
  53. return
  54. fp_write(tmp)
  55. callback(len(tmp))
  56. # return
  57. buf = ''
  58. # n = 0
  59. while True:
  60. tmp = fin.read(buf_size)
  61. # flush at EOF
  62. if not tmp:
  63. if buf:
  64. fp_write(buf)
  65. callback(1 + buf.count(delim)) # n += 1 + buf.count(delim)
  66. getattr(fout, 'flush', lambda: None)() # pragma: no cover
  67. return # n
  68. while True:
  69. try:
  70. i = tmp.index(delim)
  71. except ValueError:
  72. buf += tmp
  73. break
  74. else:
  75. fp_write(buf + tmp[:i + len(delim)])
  76. callback(1) # n += 1
  77. buf = ''
  78. tmp = tmp[i + len(delim):]
  79. # ((opt, type), ... )
  80. RE_OPTS = re.compile(r'\n {8}(\S+)\s{2,}:\s*([^,]+)')
  81. # better split method assuming no positional args
  82. RE_SHLEX = re.compile(r'\s*(?<!\S)--?([^\s=]+)(\s+|=|$)')
  83. # TODO: add custom support for some of the following?
  84. UNSUPPORTED_OPTS = ('iterable', 'gui', 'out', 'file')
  85. # The 8 leading spaces are required for consistency
  86. CLI_EXTRA_DOC = r"""
  87. Extra CLI Options
  88. -----------------
  89. name : type, optional
  90. TODO: find out why this is needed.
  91. delim : chr, optional
  92. Delimiting character [default: '\n']. Use '\0' for null.
  93. N.B.: on Windows systems, Python converts '\n' to '\r\n'.
  94. buf_size : int, optional
  95. String buffer size in bytes [default: 256]
  96. used when `delim` is specified.
  97. bytes : bool, optional
  98. If true, will count bytes, ignore `delim`, and default
  99. `unit_scale` to True, `unit_divisor` to 1024, and `unit` to 'B'.
  100. manpath : str, optional
  101. Directory in which to install tqdm man pages.
  102. comppath : str, optional
  103. Directory in which to place tqdm completion.
  104. log : str, optional
  105. CRITICAL|FATAL|ERROR|WARN(ING)|[default: 'INFO']|DEBUG|NOTSET.
  106. """
  107. def main(fp=sys.stderr, argv=None):
  108. """
  109. Parameters (internal use only)
  110. ---------
  111. fp : file-like object for tqdm
  112. argv : list (default: sys.argv[1:])
  113. """
  114. if argv is None:
  115. argv = sys.argv[1:]
  116. try:
  117. log = argv.index('--log')
  118. except ValueError:
  119. for i in argv:
  120. if i.startswith('--log='):
  121. logLevel = i[len('--log='):]
  122. break
  123. else:
  124. logLevel = 'INFO'
  125. else:
  126. # argv.pop(log)
  127. # logLevel = argv.pop(log)
  128. logLevel = argv[log + 1]
  129. logging.basicConfig(
  130. level=getattr(logging, logLevel),
  131. format="%(levelname)s:%(module)s:%(lineno)d:%(message)s")
  132. log = logging.getLogger(__name__)
  133. d = tqdm.__init__.__doc__ + CLI_EXTRA_DOC
  134. opt_types = dict(RE_OPTS.findall(d))
  135. # opt_types['delim'] = 'chr'
  136. for o in UNSUPPORTED_OPTS:
  137. opt_types.pop(o)
  138. log.debug(sorted(opt_types.items()))
  139. # d = RE_OPTS.sub(r' --\1=<\1> : \2', d)
  140. split = RE_OPTS.split(d)
  141. opt_types_desc = zip(split[1::3], split[2::3], split[3::3])
  142. d = ''.join('\n --{0}=<{0}> : {1}{2}'.format(*otd)
  143. for otd in opt_types_desc if otd[0] not in UNSUPPORTED_OPTS)
  144. d = """Usage:
  145. tqdm [--help | options]
  146. Options:
  147. -h, --help Print this help and exit
  148. -v, --version Print version and exit
  149. """ + d.strip('\n') + '\n'
  150. # opts = docopt(d, version=__version__)
  151. if any(v in argv for v in ('-v', '--version')):
  152. sys.stdout.write(__version__ + '\n')
  153. sys.exit(0)
  154. elif any(v in argv for v in ('-h', '--help')):
  155. sys.stdout.write(d + '\n')
  156. sys.exit(0)
  157. argv = RE_SHLEX.split(' '.join(["tqdm"] + argv))
  158. opts = dict(zip(argv[1::3], argv[3::3]))
  159. log.debug(opts)
  160. opts.pop('log', True)
  161. tqdm_args = {'file': fp}
  162. try:
  163. for (o, v) in opts.items():
  164. try:
  165. tqdm_args[o] = cast(v, opt_types[o])
  166. except KeyError as e:
  167. raise TqdmKeyError(str(e))
  168. log.debug('args:' + str(tqdm_args))
  169. except:
  170. fp.write('\nError:\nUsage:\n tqdm [--help | options]\n')
  171. for i in sys.stdin:
  172. sys.stdout.write(i)
  173. raise
  174. else:
  175. buf_size = tqdm_args.pop('buf_size', 256)
  176. delim = tqdm_args.pop('delim', '\n')
  177. delim_per_char = tqdm_args.pop('bytes', False)
  178. manpath = tqdm_args.pop('manpath', None)
  179. comppath = tqdm_args.pop('comppath', None)
  180. stdin = getattr(sys.stdin, 'buffer', sys.stdin)
  181. stdout = getattr(sys.stdout, 'buffer', sys.stdout)
  182. if manpath or comppath:
  183. from os import path
  184. from shutil import copyfile
  185. from pkg_resources import resource_filename, Requirement
  186. def cp(src, dst):
  187. """copies from src path to dst"""
  188. copyfile(src, dst)
  189. log.info("written:" + dst)
  190. if manpath is not None:
  191. cp(resource_filename(Requirement.parse('tqdm'), 'tqdm/tqdm.1'),
  192. path.join(manpath, 'tqdm.1'))
  193. if comppath is not None:
  194. cp(resource_filename(Requirement.parse('tqdm'),
  195. 'tqdm/completion.sh'),
  196. path.join(comppath, 'tqdm_completion.sh'))
  197. sys.exit(0)
  198. if delim_per_char:
  199. tqdm_args.setdefault('unit', 'B')
  200. tqdm_args.setdefault('unit_scale', True)
  201. tqdm_args.setdefault('unit_divisor', 1024)
  202. log.debug(tqdm_args)
  203. with tqdm(**tqdm_args) as t:
  204. posix_pipe(stdin, stdout, '', buf_size, t.update)
  205. elif delim == '\n':
  206. log.debug(tqdm_args)
  207. for i in tqdm(stdin, **tqdm_args):
  208. stdout.write(i)
  209. else:
  210. log.debug(tqdm_args)
  211. with tqdm(**tqdm_args) as t:
  212. posix_pipe(stdin, stdout, delim, buf_size, t.update)