magic_arguments.py 8.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278
  1. ''' A decorator-based method of constructing IPython magics with `argparse`
  2. option handling.
  3. New magic functions can be defined like so::
  4. from IPython.core.magic_arguments import (argument, magic_arguments,
  5. parse_argstring)
  6. @magic_arguments()
  7. @argument('-o', '--option', help='An optional argument.')
  8. @argument('arg', type=int, help='An integer positional argument.')
  9. def magic_cool(self, arg):
  10. """ A really cool magic command.
  11. """
  12. args = parse_argstring(magic_cool, arg)
  13. ...
  14. The `@magic_arguments` decorator marks the function as having argparse arguments.
  15. The `@argument` decorator adds an argument using the same syntax as argparse's
  16. `add_argument()` method. More sophisticated uses may also require the
  17. `@argument_group` or `@kwds` decorator to customize the formatting and the
  18. parsing.
  19. Help text for the magic is automatically generated from the docstring and the
  20. arguments::
  21. In[1]: %cool?
  22. %cool [-o OPTION] arg
  23. A really cool magic command.
  24. positional arguments:
  25. arg An integer positional argument.
  26. optional arguments:
  27. -o OPTION, --option OPTION
  28. An optional argument.
  29. Inheritance diagram:
  30. .. inheritance-diagram:: IPython.core.magic_arguments
  31. :parts: 3
  32. '''
  33. #-----------------------------------------------------------------------------
  34. # Copyright (C) 2010-2011, IPython Development Team.
  35. #
  36. # Distributed under the terms of the Modified BSD License.
  37. #
  38. # The full license is in the file COPYING.txt, distributed with this software.
  39. #-----------------------------------------------------------------------------
  40. import argparse
  41. import re
  42. # Our own imports
  43. from IPython.core.error import UsageError
  44. from IPython.utils.decorators import undoc
  45. from IPython.utils.process import arg_split
  46. from IPython.utils.text import dedent
  47. NAME_RE = re.compile(r"[a-zA-Z][a-zA-Z0-9_-]*$")
  48. @undoc
  49. class MagicHelpFormatter(argparse.RawDescriptionHelpFormatter):
  50. """A HelpFormatter with a couple of changes to meet our needs.
  51. """
  52. # Modified to dedent text.
  53. def _fill_text(self, text, width, indent):
  54. return argparse.RawDescriptionHelpFormatter._fill_text(self, dedent(text), width, indent)
  55. # Modified to wrap argument placeholders in <> where necessary.
  56. def _format_action_invocation(self, action):
  57. if not action.option_strings:
  58. metavar, = self._metavar_formatter(action, action.dest)(1)
  59. return metavar
  60. else:
  61. parts = []
  62. # if the Optional doesn't take a value, format is:
  63. # -s, --long
  64. if action.nargs == 0:
  65. parts.extend(action.option_strings)
  66. # if the Optional takes a value, format is:
  67. # -s ARGS, --long ARGS
  68. else:
  69. default = action.dest.upper()
  70. args_string = self._format_args(action, default)
  71. # IPYTHON MODIFICATION: If args_string is not a plain name, wrap
  72. # it in <> so it's valid RST.
  73. if not NAME_RE.match(args_string):
  74. args_string = "<%s>" % args_string
  75. for option_string in action.option_strings:
  76. parts.append('%s %s' % (option_string, args_string))
  77. return ', '.join(parts)
  78. # Override the default prefix ('usage') to our % magic escape,
  79. # in a code block.
  80. def add_usage(self, usage, actions, groups, prefix="::\n\n %"):
  81. super(MagicHelpFormatter, self).add_usage(usage, actions, groups, prefix)
  82. class MagicArgumentParser(argparse.ArgumentParser):
  83. """ An ArgumentParser tweaked for use by IPython magics.
  84. """
  85. def __init__(self,
  86. prog=None,
  87. usage=None,
  88. description=None,
  89. epilog=None,
  90. parents=None,
  91. formatter_class=MagicHelpFormatter,
  92. prefix_chars='-',
  93. argument_default=None,
  94. conflict_handler='error',
  95. add_help=False):
  96. if parents is None:
  97. parents = []
  98. super(MagicArgumentParser, self).__init__(prog=prog, usage=usage,
  99. description=description, epilog=epilog,
  100. parents=parents, formatter_class=formatter_class,
  101. prefix_chars=prefix_chars, argument_default=argument_default,
  102. conflict_handler=conflict_handler, add_help=add_help)
  103. def error(self, message):
  104. """ Raise a catchable error instead of exiting.
  105. """
  106. raise UsageError(message)
  107. def parse_argstring(self, argstring):
  108. """ Split a string into an argument list and parse that argument list.
  109. """
  110. argv = arg_split(argstring)
  111. return self.parse_args(argv)
  112. def construct_parser(magic_func):
  113. """ Construct an argument parser using the function decorations.
  114. """
  115. kwds = getattr(magic_func, 'argcmd_kwds', {})
  116. if 'description' not in kwds:
  117. kwds['description'] = getattr(magic_func, '__doc__', None)
  118. arg_name = real_name(magic_func)
  119. parser = MagicArgumentParser(arg_name, **kwds)
  120. # Reverse the list of decorators in order to apply them in the
  121. # order in which they appear in the source.
  122. group = None
  123. for deco in magic_func.decorators[::-1]:
  124. result = deco.add_to_parser(parser, group)
  125. if result is not None:
  126. group = result
  127. # Replace the magic function's docstring with the full help text.
  128. magic_func.__doc__ = parser.format_help()
  129. return parser
  130. def parse_argstring(magic_func, argstring):
  131. """ Parse the string of arguments for the given magic function.
  132. """
  133. return magic_func.parser.parse_argstring(argstring)
  134. def real_name(magic_func):
  135. """ Find the real name of the magic.
  136. """
  137. magic_name = magic_func.__name__
  138. if magic_name.startswith('magic_'):
  139. magic_name = magic_name[len('magic_'):]
  140. return getattr(magic_func, 'argcmd_name', magic_name)
  141. class ArgDecorator(object):
  142. """ Base class for decorators to add ArgumentParser information to a method.
  143. """
  144. def __call__(self, func):
  145. if not getattr(func, 'has_arguments', False):
  146. func.has_arguments = True
  147. func.decorators = []
  148. func.decorators.append(self)
  149. return func
  150. def add_to_parser(self, parser, group):
  151. """ Add this object's information to the parser, if necessary.
  152. """
  153. pass
  154. class magic_arguments(ArgDecorator):
  155. """ Mark the magic as having argparse arguments and possibly adjust the
  156. name.
  157. """
  158. def __init__(self, name=None):
  159. self.name = name
  160. def __call__(self, func):
  161. if not getattr(func, 'has_arguments', False):
  162. func.has_arguments = True
  163. func.decorators = []
  164. if self.name is not None:
  165. func.argcmd_name = self.name
  166. # This should be the first decorator in the list of decorators, thus the
  167. # last to execute. Build the parser.
  168. func.parser = construct_parser(func)
  169. return func
  170. class ArgMethodWrapper(ArgDecorator):
  171. """
  172. Base class to define a wrapper for ArgumentParser method.
  173. Child class must define either `_method_name` or `add_to_parser`.
  174. """
  175. _method_name = None
  176. def __init__(self, *args, **kwds):
  177. self.args = args
  178. self.kwds = kwds
  179. def add_to_parser(self, parser, group):
  180. """ Add this object's information to the parser.
  181. """
  182. if group is not None:
  183. parser = group
  184. getattr(parser, self._method_name)(*self.args, **self.kwds)
  185. return None
  186. class argument(ArgMethodWrapper):
  187. """ Store arguments and keywords to pass to add_argument().
  188. Instances also serve to decorate command methods.
  189. """
  190. _method_name = 'add_argument'
  191. class defaults(ArgMethodWrapper):
  192. """ Store arguments and keywords to pass to set_defaults().
  193. Instances also serve to decorate command methods.
  194. """
  195. _method_name = 'set_defaults'
  196. class argument_group(ArgMethodWrapper):
  197. """ Store arguments and keywords to pass to add_argument_group().
  198. Instances also serve to decorate command methods.
  199. """
  200. def add_to_parser(self, parser, group):
  201. """ Add this object's information to the parser.
  202. """
  203. return parser.add_argument_group(*self.args, **self.kwds)
  204. class kwds(ArgDecorator):
  205. """ Provide other keywords to the sub-parser constructor.
  206. """
  207. def __init__(self, **kwds):
  208. self.kwds = kwds
  209. def __call__(self, func):
  210. func = super(kwds, self).__call__(func)
  211. func.argcmd_kwds = self.kwds
  212. return func
  213. __all__ = ['magic_arguments', 'argument', 'argument_group', 'kwds',
  214. 'parse_argstring']