cli.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324
  1. """
  2. Module version for monitoring CLI pipes (`... | python -m tqdm | ...`).
  3. """
  4. import logging
  5. import re
  6. import sys
  7. from ast import literal_eval as numeric
  8. from textwrap import indent
  9. from .std import TqdmKeyError, TqdmTypeError, tqdm
  10. from .version import __version__
  11. __all__ = ["main"]
  12. log = logging.getLogger(__name__)
  13. def cast(val, typ):
  14. log.debug((val, typ))
  15. if " or " in typ:
  16. for t in typ.split(" or "):
  17. try:
  18. return cast(val, t)
  19. except TqdmTypeError:
  20. pass
  21. raise TqdmTypeError(f"{val} : {typ}")
  22. # sys.stderr.write('\ndebug | `val:type`: `' + val + ':' + typ + '`.\n')
  23. if typ == 'bool':
  24. if (val == 'True') or (val == ''):
  25. return True
  26. if val == 'False':
  27. return False
  28. raise TqdmTypeError(val + ' : ' + typ)
  29. if typ == 'chr':
  30. if len(val) == 1:
  31. return val.encode()
  32. if re.match(r"^\\\w+$", val):
  33. return eval(f'"{val}"').encode()
  34. raise TqdmTypeError(f"{val} : {typ}")
  35. if typ == 'str':
  36. return val
  37. if typ == 'int':
  38. try:
  39. return int(val)
  40. except ValueError as exc:
  41. raise TqdmTypeError(f"{val} : {typ}") from exc
  42. if typ == 'float':
  43. try:
  44. return float(val)
  45. except ValueError as exc:
  46. raise TqdmTypeError(f"{val} : {typ}") from exc
  47. raise TqdmTypeError(f"{val} : {typ}")
  48. def posix_pipe(fin, fout, delim=b'\\n', buf_size=256,
  49. callback=lambda float: None, callback_len=True):
  50. """
  51. Params
  52. ------
  53. fin : binary file with `read(buf_size : int)` method
  54. fout : binary file with `write` (and optionally `flush`) methods.
  55. callback : function(float), e.g.: `tqdm.update`
  56. callback_len : If (default: True) do `callback(len(buffer))`.
  57. Otherwise, do `callback(data) for data in buffer.split(delim)`.
  58. """
  59. fp_write = fout.write
  60. if not delim:
  61. while True:
  62. tmp = fin.read(buf_size)
  63. # flush at EOF
  64. if not tmp:
  65. getattr(fout, 'flush', lambda: None)()
  66. return
  67. fp_write(tmp)
  68. callback(len(tmp))
  69. # return
  70. buf = b''
  71. len_delim = len(delim)
  72. # n = 0
  73. while True:
  74. tmp = fin.read(buf_size)
  75. # flush at EOF
  76. if not tmp:
  77. if buf:
  78. fp_write(buf)
  79. if callback_len:
  80. # n += 1 + buf.count(delim)
  81. callback(1 + buf.count(delim))
  82. else:
  83. for i in buf.split(delim):
  84. callback(i)
  85. getattr(fout, 'flush', lambda: None)()
  86. return # n
  87. while True:
  88. i = tmp.find(delim)
  89. if i < 0:
  90. buf += tmp
  91. break
  92. fp_write(buf + tmp[:i + len(delim)])
  93. # n += 1
  94. callback(1 if callback_len else (buf + tmp[:i]))
  95. buf = b''
  96. tmp = tmp[i + len_delim:]
  97. # ((opt, type), ... )
  98. RE_OPTS = re.compile(r'\n {4}(\S+)\s{2,}:\s*([^,]+)')
  99. # better split method assuming no positional args
  100. RE_SHLEX = re.compile(r'\s*(?<!\S)--?([^\s=]+)(\s+|=|$)')
  101. # TODO: add custom support for some of the following?
  102. UNSUPPORTED_OPTS = ('iterable', 'gui', 'out', 'file')
  103. # The 8 leading spaces are required for consistency
  104. CLI_EXTRA_DOC = r"""
  105. Extra CLI Options
  106. -----------------
  107. name : type, optional
  108. TODO: find out why this is needed.
  109. delim : chr, optional
  110. Delimiting character [default: '\n']. Use '\0' for null.
  111. N.B.: on Windows systems, Python converts '\n' to '\r\n'.
  112. buf_size : int, optional
  113. String buffer size in bytes [default: 256]
  114. used when `delim` is specified.
  115. bytes : bool, optional
  116. If true, will count bytes, ignore `delim`, and default
  117. `unit_scale` to True, `unit_divisor` to 1024, and `unit` to 'B'.
  118. tee : bool, optional
  119. If true, passes `stdin` to both `stderr` and `stdout`.
  120. update : bool, optional
  121. If true, will treat input as newly elapsed iterations,
  122. i.e. numbers to pass to `update()`. Note that this is slow
  123. (~2e5 it/s) since every input must be decoded as a number.
  124. update_to : bool, optional
  125. If true, will treat input as total elapsed iterations,
  126. i.e. numbers to assign to `self.n`. Note that this is slow
  127. (~2e5 it/s) since every input must be decoded as a number.
  128. null : bool, optional
  129. If true, will discard input (no stdout).
  130. manpath : str, optional
  131. Directory in which to install tqdm man pages.
  132. comppath : str, optional
  133. Directory in which to place tqdm completion.
  134. log : str, optional
  135. CRITICAL|FATAL|ERROR|WARN(ING)|[default: 'INFO']|DEBUG|NOTSET.
  136. """
  137. def main(fp=sys.stderr, argv=None):
  138. """
  139. Parameters (internal use only)
  140. ---------
  141. fp : file-like object for tqdm
  142. argv : list (default: sys.argv[1:])
  143. """
  144. if argv is None:
  145. argv = sys.argv[1:]
  146. try:
  147. log_idx = argv.index('--log')
  148. except ValueError:
  149. for i in argv:
  150. if i.startswith('--log='):
  151. logLevel = i[len('--log='):]
  152. break
  153. else:
  154. logLevel = 'INFO'
  155. else:
  156. # argv.pop(log_idx)
  157. # logLevel = argv.pop(log_idx)
  158. logLevel = argv[log_idx + 1]
  159. logging.basicConfig(level=getattr(logging, logLevel),
  160. format="%(levelname)s:%(module)s:%(lineno)d:%(message)s")
  161. # py<3.13 doesn't dedent docstrings
  162. d = (tqdm.__doc__ if sys.version_info < (3, 13)
  163. else indent(tqdm.__doc__, " ")) + CLI_EXTRA_DOC
  164. opt_types = dict(RE_OPTS.findall(d))
  165. # opt_types['delim'] = 'chr'
  166. for o in UNSUPPORTED_OPTS:
  167. opt_types.pop(o)
  168. log.debug(sorted(opt_types.items()))
  169. # d = RE_OPTS.sub(r' --\1=<\1> : \2', d)
  170. split = RE_OPTS.split(d)
  171. opt_types_desc = zip(split[1::3], split[2::3], split[3::3])
  172. d = ''.join(('\n --{0} : {2}{3}' if otd[1] == 'bool' else
  173. '\n --{0}=<{1}> : {2}{3}').format(
  174. otd[0].replace('_', '-'), otd[0], *otd[1:])
  175. for otd in opt_types_desc if otd[0] not in UNSUPPORTED_OPTS)
  176. help_short = "Usage:\n tqdm [--help | options]\n"
  177. d = help_short + """
  178. Options:
  179. -h, --help Print this help and exit.
  180. -v, --version Print version and exit.
  181. """ + d.strip('\n') + '\n'
  182. # opts = docopt(d, version=__version__)
  183. if any(v in argv for v in ('-v', '--version')):
  184. sys.stdout.write(__version__ + '\n')
  185. sys.exit(0)
  186. elif any(v in argv for v in ('-h', '--help')):
  187. sys.stdout.write(d + '\n')
  188. sys.exit(0)
  189. elif argv and argv[0][:2] != '--':
  190. sys.stderr.write(f"Error:Unknown argument:{argv[0]}\n{help_short}")
  191. argv = RE_SHLEX.split(' '.join(["tqdm"] + argv))
  192. opts = dict(zip(argv[1::3], argv[3::3]))
  193. log.debug(opts)
  194. opts.pop('log', True)
  195. tqdm_args = {'file': fp}
  196. try:
  197. for (o, v) in opts.items():
  198. o = o.replace('-', '_')
  199. try:
  200. tqdm_args[o] = cast(v, opt_types[o])
  201. except KeyError as e:
  202. raise TqdmKeyError(str(e))
  203. log.debug('args:' + str(tqdm_args))
  204. delim_per_char = tqdm_args.pop('bytes', False)
  205. update = tqdm_args.pop('update', False)
  206. update_to = tqdm_args.pop('update_to', False)
  207. if sum((delim_per_char, update, update_to)) > 1:
  208. raise TqdmKeyError("Can only have one of --bytes --update --update_to")
  209. except Exception:
  210. fp.write("\nError:\n" + help_short)
  211. stdin, stdout_write = sys.stdin, sys.stdout.write
  212. for i in stdin:
  213. stdout_write(i)
  214. raise
  215. else:
  216. buf_size = tqdm_args.pop('buf_size', 256)
  217. delim = tqdm_args.pop('delim', b'\\n')
  218. tee = tqdm_args.pop('tee', False)
  219. manpath = tqdm_args.pop('manpath', None)
  220. comppath = tqdm_args.pop('comppath', None)
  221. if tqdm_args.pop('null', False):
  222. class stdout(object):
  223. @staticmethod
  224. def write(_):
  225. pass
  226. else:
  227. stdout = sys.stdout
  228. stdout = getattr(stdout, 'buffer', stdout)
  229. stdin = getattr(sys.stdin, 'buffer', sys.stdin)
  230. if manpath or comppath:
  231. try: # py<3.9
  232. import importlib_resources as resources
  233. except ImportError:
  234. from importlib import resources
  235. from pathlib import Path
  236. def cp(name, dst):
  237. """copy resource `name` to `dst`"""
  238. fi = resources.files('tqdm') / name
  239. dst.write_bytes(fi.read_bytes())
  240. log.info("written:%s", dst)
  241. if manpath is not None:
  242. cp('tqdm.1', Path(manpath) / 'tqdm.1')
  243. if comppath is not None:
  244. cp('completion.sh', Path(comppath) / 'tqdm_completion.sh')
  245. sys.exit(0)
  246. if tee:
  247. stdout_write = stdout.write
  248. fp_write = getattr(fp, 'buffer', fp).write
  249. class stdout(object): # pylint: disable=function-redefined
  250. @staticmethod
  251. def write(x):
  252. with tqdm.external_write_mode(file=fp):
  253. fp_write(x)
  254. stdout_write(x)
  255. if delim_per_char:
  256. tqdm_args.setdefault('unit', 'B')
  257. tqdm_args.setdefault('unit_scale', True)
  258. tqdm_args.setdefault('unit_divisor', 1024)
  259. log.debug(tqdm_args)
  260. with tqdm(**tqdm_args) as t:
  261. posix_pipe(stdin, stdout, '', buf_size, t.update)
  262. elif delim == b'\\n':
  263. log.debug(tqdm_args)
  264. write = stdout.write
  265. if update or update_to:
  266. with tqdm(**tqdm_args) as t:
  267. if update:
  268. def callback(i):
  269. t.update(numeric(i.decode()))
  270. else: # update_to
  271. def callback(i):
  272. t.update(numeric(i.decode()) - t.n)
  273. for i in stdin:
  274. write(i)
  275. callback(i)
  276. else:
  277. for i in tqdm(stdin, **tqdm_args):
  278. write(i)
  279. else:
  280. log.debug(tqdm_args)
  281. with tqdm(**tqdm_args) as t:
  282. callback_len = False
  283. if update:
  284. def callback(i):
  285. t.update(numeric(i.decode()))
  286. elif update_to:
  287. def callback(i):
  288. t.update(numeric(i.decode()) - t.n)
  289. else:
  290. callback = t.update
  291. callback_len = True
  292. posix_pipe(stdin, stdout, delim, buf_size, callback, callback_len)