You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

382 lines
12 KiB

4 years ago
  1. # coding: utf-8
  2. #
  3. # Copyright © 2010—2014 Andrey Mikhaylenko and contributors
  4. #
  5. # This file is part of Argh.
  6. #
  7. # Argh is free software under terms of the GNU Lesser
  8. # General Public License version 3 (LGPLv3) as published by the Free
  9. # Software Foundation. See the file README.rst for copying conditions.
  10. #
  11. """
  12. Dispatching
  13. ~~~~~~~~~~~
  14. """
  15. import argparse
  16. import sys
  17. from types import GeneratorType
  18. from argh import compat, io
  19. from argh.constants import (
  20. ATTR_WRAPPED_EXCEPTIONS,
  21. ATTR_WRAPPED_EXCEPTIONS_PROCESSOR,
  22. ATTR_EXPECTS_NAMESPACE_OBJECT,
  23. PARSER_FORMATTER,
  24. DEST_FUNCTION,
  25. )
  26. from argh.completion import autocomplete
  27. from argh.assembling import add_commands, set_default_command
  28. from argh.exceptions import DispatchingError, CommandError
  29. from argh.utils import get_arg_spec
  30. __all__ = ['dispatch', 'dispatch_command', 'dispatch_commands',
  31. 'PARSER_FORMATTER', 'EntryPoint']
  32. class ArghNamespace(argparse.Namespace):
  33. """
  34. A namespace object which collects the stack of functions (the
  35. :attr:`~argh.constants.DEST_FUNCTION` arguments passed to it via
  36. parser's defaults).
  37. """
  38. def __init__(self, *args, **kw):
  39. super(ArghNamespace, self).__init__(*args, **kw)
  40. self._functions_stack = []
  41. def __setattr__(self, k, v):
  42. if k == DEST_FUNCTION:
  43. # don't register the function under DEST_FUNCTION name.
  44. # If `ArgumentParser.parse_known_args()` sees that we already have
  45. # such attribute, it skips it. However, it goes from the topmost
  46. # parser to subparsers. We need the function mapped to the
  47. # subparser. So we fool the `ArgumentParser` and pretend that we
  48. # didn't get a DEST_FUNCTION attribute; however, in fact we collect
  49. # all its values in a stack. The last item in the stack would be
  50. # the function mapped to the innermost parser — the one we need.
  51. self._functions_stack.append(v)
  52. else:
  53. super(ArghNamespace, self).__setattr__(k, v)
  54. def get_function(self):
  55. return self._functions_stack[-1]
  56. def dispatch(parser, argv=None, add_help_command=True,
  57. completion=True, pre_call=None,
  58. output_file=sys.stdout, errors_file=sys.stderr,
  59. raw_output=False, namespace=None,
  60. skip_unknown_args=False):
  61. """
  62. Parses given list of arguments using given parser, calls the relevant
  63. function and prints the result.
  64. The target function should expect one positional argument: the
  65. :class:`argparse.Namespace` object. However, if the function is decorated with
  66. :func:`~argh.decorators.plain_signature`, the positional and named
  67. arguments from the namespace object are passed to the function instead
  68. of the object itself.
  69. :param parser:
  70. the ArgumentParser instance.
  71. :param argv:
  72. a list of strings representing the arguments. If `None`, ``sys.argv``
  73. is used instead. Default is `None`.
  74. :param add_help_command:
  75. if `True`, converts first positional argument "help" to a keyword
  76. argument so that ``help foo`` becomes ``foo --help`` and displays usage
  77. information for "foo". Default is `True`.
  78. :param output_file:
  79. A file-like object for output. If `None`, the resulting lines are
  80. collected and returned as a string. Default is ``sys.stdout``.
  81. :param errors_file:
  82. Same as `output_file` but for ``sys.stderr``.
  83. :param raw_output:
  84. If `True`, results are written to the output file raw, without adding
  85. whitespaces or newlines between yielded strings. Default is `False`.
  86. :param completion:
  87. If `True`, shell tab completion is enabled. Default is `True`. (You
  88. will also need to install it.) See :mod:`argh.completion`.
  89. :param skip_unknown_args:
  90. If `True`, unknown arguments do not cause an error
  91. (`ArgumentParser.parse_known_args` is used).
  92. :param namespace:
  93. An `argparse.Namespace`-like object. By default an
  94. :class:`ArghNamespace` object is used. Please note that support for
  95. combined default and nested functions may be broken if a different
  96. type of object is forced.
  97. By default the exceptions are not wrapped and will propagate. The only
  98. exception that is always wrapped is :class:`~argh.exceptions.CommandError`
  99. which is interpreted as an expected event so the traceback is hidden.
  100. You can also mark arbitrary exceptions as "wrappable" by using the
  101. :func:`~argh.decorators.wrap_errors` decorator.
  102. """
  103. if completion:
  104. autocomplete(parser)
  105. if argv is None:
  106. argv = sys.argv[1:]
  107. if add_help_command:
  108. if argv and argv[0] == 'help':
  109. argv.pop(0)
  110. argv.append('--help')
  111. if skip_unknown_args:
  112. parse_args = parser.parse_known_args
  113. else:
  114. parse_args = parser.parse_args
  115. if not namespace:
  116. namespace = ArghNamespace()
  117. # this will raise SystemExit if parsing fails
  118. namespace_obj = parse_args(argv, namespace=namespace)
  119. function = _get_function_from_namespace_obj(namespace_obj)
  120. if function:
  121. lines = _execute_command(function, namespace_obj, errors_file,
  122. pre_call=pre_call)
  123. else:
  124. # no commands declared, can't dispatch; display help message
  125. lines = [parser.format_usage()]
  126. if output_file is None:
  127. # user wants a string; we create an internal temporary file-like object
  128. # and will return its contents as a string
  129. if sys.version_info < (3,0):
  130. f = compat.BytesIO()
  131. else:
  132. f = compat.StringIO()
  133. else:
  134. # normally this is stdout; can be any file
  135. f = output_file
  136. for line in lines:
  137. # print the line as soon as it is generated to ensure that it is
  138. # displayed to the user before anything else happens, e.g.
  139. # raw_input() is called
  140. io.dump(line, f)
  141. if not raw_output:
  142. # in most cases user wants one message per line
  143. io.dump('\n', f)
  144. if output_file is None:
  145. # user wanted a string; return contents of our temporary file-like obj
  146. f.seek(0)
  147. return f.read()
  148. def _get_function_from_namespace_obj(namespace_obj):
  149. if isinstance(namespace_obj, ArghNamespace):
  150. # our special namespace object keeps the stack of assigned functions
  151. try:
  152. function = namespace_obj.get_function()
  153. except (AttributeError, IndexError):
  154. return None
  155. else:
  156. # a custom (probably vanilla) namespace object keeps the last assigned
  157. # function; this may be wrong but at least something may work
  158. if not hasattr(namespace_obj, DEST_FUNCTION):
  159. return None
  160. function = getattr(namespace_obj, DEST_FUNCTION)
  161. if not function or not hasattr(function, '__call__'):
  162. return None
  163. return function
  164. def _execute_command(function, namespace_obj, errors_file, pre_call=None):
  165. """
  166. Assumes that `function` is a callable. Tries different approaches
  167. to call it (with `namespace_obj` or with ordinary signature).
  168. Yields the results line by line.
  169. If :class:`~argh.exceptions.CommandError` is raised, its message is
  170. appended to the results (i.e. yielded by the generator as a string).
  171. All other exceptions propagate unless marked as wrappable
  172. by :func:`wrap_errors`.
  173. """
  174. if pre_call: # XXX undocumented because I'm unsure if it's OK
  175. # Actually used in real projects:
  176. # * https://google.com/search?q=argh+dispatch+pre_call
  177. # * https://github.com/neithere/argh/issues/63
  178. pre_call(namespace_obj)
  179. # the function is nested to catch certain exceptions (see below)
  180. def _call():
  181. # Actually call the function
  182. if getattr(function, ATTR_EXPECTS_NAMESPACE_OBJECT, False):
  183. result = function(namespace_obj)
  184. else:
  185. # namespace -> dictionary
  186. _flat_key = lambda key: key.replace('-', '_')
  187. all_input = dict((_flat_key(k), v)
  188. for k,v in vars(namespace_obj).items())
  189. # filter the namespace variables so that only those expected
  190. # by the actual function will pass
  191. spec = get_arg_spec(function)
  192. positional = [all_input[k] for k in spec.args]
  193. kwonly = getattr(spec, 'kwonlyargs', [])
  194. keywords = dict((k, all_input[k]) for k in kwonly)
  195. # *args
  196. if spec.varargs:
  197. positional += getattr(namespace_obj, spec.varargs)
  198. # **kwargs
  199. varkw = getattr(spec, 'varkw', getattr(spec, 'keywords', []))
  200. if varkw:
  201. not_kwargs = [DEST_FUNCTION] + spec.args + [spec.varargs] + kwonly
  202. for k in vars(namespace_obj):
  203. if k.startswith('_') or k in not_kwargs:
  204. continue
  205. keywords[k] = getattr(namespace_obj, k)
  206. result = function(*positional, **keywords)
  207. # Yield the results
  208. if isinstance(result, (GeneratorType, list, tuple)):
  209. # yield each line ASAP, convert CommandError message to a line
  210. for line in result:
  211. yield line
  212. else:
  213. # yield non-empty non-iterable result as a single line
  214. if result is not None:
  215. yield result
  216. wrappable_exceptions = [CommandError]
  217. wrappable_exceptions += getattr(function, ATTR_WRAPPED_EXCEPTIONS, [])
  218. try:
  219. result = _call()
  220. for line in result:
  221. yield line
  222. except tuple(wrappable_exceptions) as e:
  223. processor = getattr(function, ATTR_WRAPPED_EXCEPTIONS_PROCESSOR,
  224. lambda e: '{0.__class__.__name__}: {0}'.format(e))
  225. errors_file.write(compat.text_type(processor(e)))
  226. errors_file.write('\n')
  227. def dispatch_command(function, *args, **kwargs):
  228. """
  229. A wrapper for :func:`dispatch` that creates a one-command parser.
  230. Uses :attr:`PARSER_FORMATTER`.
  231. This::
  232. dispatch_command(foo)
  233. ...is a shortcut for::
  234. parser = ArgumentParser()
  235. set_default_command(parser, foo)
  236. dispatch(parser)
  237. This function can be also used as a decorator.
  238. """
  239. parser = argparse.ArgumentParser(formatter_class=PARSER_FORMATTER)
  240. set_default_command(parser, function)
  241. dispatch(parser, *args, **kwargs)
  242. def dispatch_commands(functions, *args, **kwargs):
  243. """
  244. A wrapper for :func:`dispatch` that creates a parser, adds commands to
  245. the parser and dispatches them.
  246. Uses :attr:`PARSER_FORMATTER`.
  247. This::
  248. dispatch_commands([foo, bar])
  249. ...is a shortcut for::
  250. parser = ArgumentParser()
  251. add_commands(parser, [foo, bar])
  252. dispatch(parser)
  253. """
  254. parser = argparse.ArgumentParser(formatter_class=PARSER_FORMATTER)
  255. add_commands(parser, functions)
  256. dispatch(parser, *args, **kwargs)
  257. class EntryPoint(object):
  258. """
  259. An object to which functions can be attached and then dispatched.
  260. When called with an argument, the argument (a function) is registered
  261. at this entry point as a command.
  262. When called without an argument, dispatching is triggered with all
  263. previously registered commands.
  264. Usage::
  265. from argh import EntryPoint
  266. app = EntryPoint('main', dict(description='This is a cool app'))
  267. @app
  268. def ls():
  269. for i in range(10):
  270. print i
  271. @app
  272. def greet():
  273. print 'hello'
  274. if __name__ == '__main__':
  275. app()
  276. """
  277. def __init__(self, name=None, parser_kwargs=None):
  278. self.name = name or 'unnamed'
  279. self.commands = []
  280. self.parser_kwargs = parser_kwargs or {}
  281. def __call__(self, f=None):
  282. if f:
  283. self._register_command(f)
  284. return f
  285. return self._dispatch()
  286. def _register_command(self, f):
  287. self.commands.append(f)
  288. def _dispatch(self):
  289. if not self.commands:
  290. raise DispatchingError('no commands for entry point "{0}"'
  291. .format(self.name))
  292. parser = argparse.ArgumentParser(**self.parser_kwargs)
  293. add_commands(parser, self.commands)
  294. dispatch(parser)