__init__.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416
  1. # -*- coding: utf-8 -*-
  2. from __future__ import absolute_import
  3. import os
  4. import re
  5. import sys
  6. import types
  7. import warnings
  8. from gettext import gettext as _
  9. import argparse
  10. from flask import Flask
  11. from ._compat import iteritems
  12. from .commands import Group, Option, Command, Server, Shell
  13. from .cli import prompt, prompt_pass, prompt_bool, prompt_choices
  14. __all__ = ["Command", "Shell", "Server", "Manager", "Group", "Option",
  15. "prompt", "prompt_pass", "prompt_bool", "prompt_choices"]
  16. safe_actions = (argparse._StoreAction,
  17. argparse._StoreConstAction,
  18. argparse._StoreTrueAction,
  19. argparse._StoreFalseAction,
  20. argparse._AppendAction,
  21. argparse._AppendConstAction,
  22. argparse._CountAction)
  23. try:
  24. import argcomplete
  25. ARGCOMPLETE_IMPORTED = True
  26. except ImportError:
  27. ARGCOMPLETE_IMPORTED = False
  28. def add_help(parser, help_args):
  29. if not help_args:
  30. return
  31. parser.add_argument(*help_args,
  32. action='help', default=argparse.SUPPRESS, help=_('show this help message and exit'))
  33. class Manager(object):
  34. """
  35. Controller class for handling a set of commands.
  36. Typical usage::
  37. class Print(Command):
  38. def run(self):
  39. print "hello"
  40. app = Flask(__name__)
  41. manager = Manager(app)
  42. manager.add_command("print", Print())
  43. if __name__ == "__main__":
  44. manager.run()
  45. On command line::
  46. python manage.py print
  47. > hello
  48. :param app: Flask instance, or callable returning a Flask instance.
  49. :param with_default_commands: load commands **runserver** and **shell**
  50. by default.
  51. :param disable_argcomplete: disable automatic loading of argcomplete.
  52. """
  53. help_args = ('-?','--help')
  54. def __init__(self, app=None, with_default_commands=None, usage=None,
  55. help=None, description=None, disable_argcomplete=False):
  56. self.app = app
  57. self._commands = dict()
  58. self._options = list()
  59. self.usage = usage
  60. self.help = help if help is not None else usage
  61. self.description = description if description is not None else usage
  62. self.disable_argcomplete = disable_argcomplete
  63. self.with_default_commands = with_default_commands
  64. self.parent = None
  65. def add_default_commands(self):
  66. """
  67. Adds the shell and runserver default commands. To override these,
  68. simply add your own equivalents using add_command or decorators.
  69. """
  70. if "shell" not in self._commands:
  71. self.add_command("shell", Shell())
  72. if "runserver" not in self._commands:
  73. self.add_command("runserver", Server())
  74. def add_option(self, *args, **kwargs):
  75. """
  76. Adds a global option. This is useful if you want to set variables
  77. applying to the application setup, rather than individual commands.
  78. For this to work, the manager must be initialized with a factory
  79. function rather than a Flask instance. Otherwise any options you set
  80. will be ignored.
  81. The arguments are then passed to your function, e.g.::
  82. def create_my_app(config=None):
  83. app = Flask(__name__)
  84. if config:
  85. app.config.from_pyfile(config)
  86. return app
  87. manager = Manager(create_my_app)
  88. manager.add_option("-c", "--config", dest="config", required=False)
  89. @manager.command
  90. def mycommand(app):
  91. app.do_something()
  92. and are invoked like this::
  93. > python manage.py -c dev.cfg mycommand
  94. Any manager options passed on the command line will not be passed to
  95. the command.
  96. Arguments for this function are the same as for the Option class.
  97. """
  98. self._options.append(Option(*args, **kwargs))
  99. def __call__(self, app=None, **kwargs):
  100. """
  101. This procedure is called with the App instance (if this is a
  102. sub-Manager) and any options.
  103. If your sub-Manager does not override this, any values for options will get lost.
  104. """
  105. if app is None:
  106. app = self.app
  107. if app is None:
  108. raise Exception("There is no app here. This is unlikely to work.")
  109. if isinstance(app, Flask):
  110. if kwargs:
  111. import warnings
  112. warnings.warn("Options will be ignored.")
  113. return app
  114. app = app(**kwargs)
  115. self.app = app
  116. return app
  117. def create_app(self, *args, **kwargs):
  118. warnings.warn("create_app() is deprecated; use __call__().", warnings.DeprecationWarning)
  119. return self(*args,**kwargs)
  120. def create_parser(self, prog, func_stack=(), parent=None):
  121. """
  122. Creates an ArgumentParser instance from options returned
  123. by get_options(), and subparser for the given commands.
  124. """
  125. prog = os.path.basename(prog)
  126. func_stack=func_stack+(self,)
  127. options_parser = argparse.ArgumentParser(add_help=False)
  128. for option in self.get_options():
  129. options_parser.add_argument(*option.args, **option.kwargs)
  130. parser = argparse.ArgumentParser(prog=prog, usage=self.usage,
  131. description=self.description,
  132. parents=[options_parser],
  133. add_help=False)
  134. add_help(parser, self.help_args)
  135. self._patch_argparser(parser)
  136. subparsers = parser.add_subparsers()
  137. for name, command in self._commands.items():
  138. usage = getattr(command, 'usage', None)
  139. help = getattr(command, 'help', None)
  140. if help is None: help = command.__doc__
  141. description = getattr(command, 'description', None)
  142. if description is None: description = command.__doc__
  143. command_parser = command.create_parser(name, func_stack=func_stack, parent=self)
  144. subparser = subparsers.add_parser(name, usage=usage, help=help,
  145. description=description,
  146. parents=[command_parser],
  147. add_help=False)
  148. if isinstance(command, Manager):
  149. self._patch_argparser(subparser)
  150. ## enable autocomplete only for parent parser when argcomplete is
  151. ## imported and it is NOT disabled in constructor
  152. if parent is None and ARGCOMPLETE_IMPORTED \
  153. and not self.disable_argcomplete:
  154. argcomplete.autocomplete(parser, always_complete_options=True)
  155. self.parser = parser
  156. return parser
  157. # def foo(self, app, *args, **kwargs):
  158. # print(args)
  159. def _patch_argparser(self, parser):
  160. """
  161. Patches the parser to print the full help if no arguments are supplied
  162. """
  163. def _parse_known_args(self, arg_strings, *args, **kw):
  164. if not arg_strings:
  165. self.print_help()
  166. self.exit(2)
  167. return self._parse_known_args2(arg_strings, *args, **kw)
  168. parser._parse_known_args2 = parser._parse_known_args
  169. parser._parse_known_args = types.MethodType(_parse_known_args, parser)
  170. def get_options(self):
  171. return self._options
  172. def add_command(self, *args, **kwargs):
  173. """
  174. Adds command to registry.
  175. :param command: Command instance
  176. :param name: Name of the command (optional)
  177. :param namespace: Namespace of the command (optional; pass as kwarg)
  178. """
  179. if len(args) == 1:
  180. command = args[0]
  181. name = None
  182. else:
  183. name, command = args
  184. if name is None:
  185. if hasattr(command, 'name'):
  186. name = command.name
  187. else:
  188. name = type(command).__name__.lower()
  189. name = re.sub(r'command$', '', name)
  190. if isinstance(command, Manager):
  191. command.parent = self
  192. if isinstance(command, type):
  193. command = command()
  194. namespace = kwargs.get('namespace')
  195. if not namespace:
  196. namespace = getattr(command, 'namespace', None)
  197. if namespace:
  198. if namespace not in self._commands:
  199. self.add_command(namespace, Manager())
  200. self._commands[namespace]._commands[name] = command
  201. else:
  202. self._commands[name] = command
  203. def command(self, func):
  204. """
  205. Decorator to add a command function to the registry.
  206. :param func: command function.Arguments depend on the
  207. options.
  208. """
  209. command = Command(func)
  210. self.add_command(func.__name__, command)
  211. return func
  212. def option(self, *args, **kwargs):
  213. """
  214. Decorator to add an option to a function. Automatically registers the
  215. function - do not use together with ``@command``. You can add as many
  216. ``@option`` calls as you like, for example::
  217. @option('-n', '--name', dest='name')
  218. @option('-u', '--url', dest='url')
  219. def hello(name, url):
  220. print "hello", name, url
  221. Takes the same arguments as the ``Option`` constructor.
  222. """
  223. option = Option(*args, **kwargs)
  224. def decorate(func):
  225. name = func.__name__
  226. if name not in self._commands:
  227. command = Command()
  228. command.run = func
  229. command.__doc__ = func.__doc__
  230. command.option_list = []
  231. self.add_command(name, command)
  232. self._commands[name].option_list.append(option)
  233. return func
  234. return decorate
  235. def shell(self, func):
  236. """
  237. Decorator that wraps function in shell command. This is equivalent to::
  238. def _make_context(app):
  239. return dict(app=app)
  240. manager.add_command("shell", Shell(make_context=_make_context))
  241. The decorated function should take a single "app" argument, and return
  242. a dict.
  243. For more sophisticated usage use the Shell class.
  244. """
  245. self.add_command('shell', Shell(make_context=func))
  246. return func
  247. def set_defaults(self):
  248. if self.with_default_commands is None:
  249. self.with_default_commands = self.parent is None
  250. if self.with_default_commands:
  251. self.add_default_commands()
  252. self.with_default_commands = False
  253. def handle(self, prog, args=None):
  254. self.set_defaults()
  255. app_parser = self.create_parser(prog)
  256. args = list(args or [])
  257. app_namespace, remaining_args = app_parser.parse_known_args(args)
  258. # get the handle function and remove it from parsed options
  259. kwargs = app_namespace.__dict__
  260. func_stack = kwargs.pop('func_stack', None)
  261. if not func_stack:
  262. app_parser.error('too few arguments')
  263. last_func = func_stack[-1]
  264. if remaining_args and not getattr(last_func, 'capture_all_args', False):
  265. app_parser.error('too many arguments')
  266. args = []
  267. for handle in func_stack:
  268. # get only safe config options
  269. config_keys = [action.dest for action in handle.parser._actions
  270. if handle is last_func or action.__class__ in safe_actions]
  271. # pass only safe app config keys
  272. config = dict((k, v) for k, v in iteritems(kwargs)
  273. if k in config_keys)
  274. # remove application config keys from handle kwargs
  275. kwargs = dict((k, v) for k, v in iteritems(kwargs)
  276. if k not in config_keys)
  277. if handle is last_func and getattr(last_func, 'capture_all_args', False):
  278. args.append(remaining_args)
  279. try:
  280. res = handle(*args, **config)
  281. except TypeError as err:
  282. err.args = ("{}: {}".format(handle,str(err)),)
  283. raise
  284. args = [res]
  285. assert not kwargs
  286. return res
  287. def run(self, commands=None, default_command=None):
  288. """
  289. Prepares manager to receive command line input. Usually run
  290. inside "if __name__ == "__main__" block in a Python script.
  291. :param commands: optional dict of commands. Appended to any commands
  292. added using add_command().
  293. :param default_command: name of default command to run if no
  294. arguments passed.
  295. """
  296. if commands:
  297. self._commands.update(commands)
  298. if default_command is not None and len(sys.argv) == 1:
  299. sys.argv.append(default_command)
  300. try:
  301. result = self.handle(sys.argv[0], sys.argv[1:])
  302. except SystemExit as e:
  303. result = e.code
  304. sys.exit(result or 0)