_shellcomp.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667
  1. # -*- test-case-name: twisted.python.test.test_shellcomp -*-
  2. # Copyright (c) Twisted Matrix Laboratories.
  3. # See LICENSE for details.
  4. """
  5. No public APIs are provided by this module. Internal use only.
  6. This module implements dynamic tab-completion for any command that uses
  7. twisted.python.usage. Currently, only zsh is supported. Bash support may
  8. be added in the future.
  9. Maintainer: Eric P. Mangold - twisted AT teratorn DOT org
  10. In order for zsh completion to take place the shell must be able to find an
  11. appropriate "stub" file ("completion function") that invokes this code and
  12. displays the results to the user.
  13. The stub used for Twisted commands is in the file C{twisted-completion.zsh},
  14. which is also included in the official Zsh distribution at
  15. C{Completion/Unix/Command/_twisted}. Use this file as a basis for completion
  16. functions for your own commands. You should only need to change the first line
  17. to something like C{#compdef mycommand}.
  18. The main public documentation exists in the L{twisted.python.usage.Options}
  19. docstring, the L{twisted.python.usage.Completions} docstring, and the
  20. Options howto.
  21. """
  22. import itertools, getopt, inspect
  23. from twisted.python import reflect, util, usage
  24. def shellComplete(config, cmdName, words, shellCompFile):
  25. """
  26. Perform shell completion.
  27. A completion function (shell script) is generated for the requested
  28. shell and written to C{shellCompFile}, typically C{stdout}. The result
  29. is then eval'd by the shell to produce the desired completions.
  30. @type config: L{twisted.python.usage.Options}
  31. @param config: The L{twisted.python.usage.Options} instance to generate
  32. completions for.
  33. @type cmdName: C{str}
  34. @param cmdName: The name of the command we're generating completions for.
  35. In the case of zsh, this is used to print an appropriate
  36. "#compdef $CMD" line at the top of the output. This is
  37. not necessary for the functionality of the system, but it
  38. helps in debugging, since the output we produce is properly
  39. formed and may be saved in a file and used as a stand-alone
  40. completion function.
  41. @type words: C{list} of C{str}
  42. @param words: The raw command-line words passed to use by the shell
  43. stub function. argv[0] has already been stripped off.
  44. @type shellCompFile: C{file}
  45. @param shellCompFile: The file to write completion data to.
  46. """
  47. # shellName is provided for forward-compatibility. It is not used,
  48. # since we currently only support zsh.
  49. shellName, position = words[-1].split(":")
  50. position = int(position)
  51. # zsh gives the completion position ($CURRENT) as a 1-based index,
  52. # and argv[0] has already been stripped off, so we subtract 2 to
  53. # get the real 0-based index.
  54. position -= 2
  55. cWord = words[position]
  56. # since the user may hit TAB at any time, we may have been called with an
  57. # incomplete command-line that would generate getopt errors if parsed
  58. # verbatim. However, we must do *some* parsing in order to determine if
  59. # there is a specific subcommand that we need to provide completion for.
  60. # So, to make the command-line more sane we work backwards from the
  61. # current completion position and strip off all words until we find one
  62. # that "looks" like a subcommand. It may in fact be the argument to a
  63. # normal command-line option, but that won't matter for our purposes.
  64. while position >= 1:
  65. if words[position - 1].startswith("-"):
  66. position -= 1
  67. else:
  68. break
  69. words = words[:position]
  70. subCommands = getattr(config, 'subCommands', None)
  71. if subCommands:
  72. # OK, this command supports sub-commands, so lets see if we have been
  73. # given one.
  74. # If the command-line arguments are not valid then we won't be able to
  75. # sanely detect the sub-command, so just generate completions as if no
  76. # sub-command was found.
  77. args = None
  78. try:
  79. opts, args = getopt.getopt(words,
  80. config.shortOpt, config.longOpt)
  81. except getopt.error:
  82. pass
  83. if args:
  84. # yes, we have a subcommand. Try to find it.
  85. for (cmd, short, parser, doc) in config.subCommands:
  86. if args[0] == cmd or args[0] == short:
  87. subOptions = parser()
  88. subOptions.parent = config
  89. gen = ZshSubcommandBuilder(subOptions, config, cmdName,
  90. shellCompFile)
  91. gen.write()
  92. return
  93. # sub-command not given, or did not match any knowns sub-command names
  94. genSubs = True
  95. if cWord.startswith("-"):
  96. # optimization: if the current word being completed starts
  97. # with a hyphen then it can't be a sub-command, so skip
  98. # the expensive generation of the sub-command list
  99. genSubs = False
  100. gen = ZshBuilder(config, cmdName, shellCompFile)
  101. gen.write(genSubs=genSubs)
  102. else:
  103. gen = ZshBuilder(config, cmdName, shellCompFile)
  104. gen.write()
  105. class SubcommandAction(usage.Completer):
  106. def _shellCode(self, optName, shellType):
  107. if shellType == usage._ZSH:
  108. return '*::subcmd:->subcmd'
  109. raise NotImplementedError("Unknown shellType %r" % (shellType,))
  110. class ZshBuilder(object):
  111. """
  112. Constructs zsh code that will complete options for a given usage.Options
  113. instance, possibly including a list of subcommand names.
  114. Completions for options to subcommands won't be generated because this
  115. class will never be used if the user is completing options for a specific
  116. subcommand. (See L{ZshSubcommandBuilder} below)
  117. @type options: L{twisted.python.usage.Options}
  118. @ivar options: The L{twisted.python.usage.Options} instance defined for this
  119. command.
  120. @type cmdName: C{str}
  121. @ivar cmdName: The name of the command we're generating completions for.
  122. @type file: C{file}
  123. @ivar file: The C{file} to write the completion function to.
  124. """
  125. def __init__(self, options, cmdName, file):
  126. self.options = options
  127. self.cmdName = cmdName
  128. self.file = file
  129. def write(self, genSubs=True):
  130. """
  131. Generate the completion function and write it to the output file
  132. @return: L{None}
  133. @type genSubs: C{bool}
  134. @param genSubs: Flag indicating whether or not completions for the list
  135. of subcommand should be generated. Only has an effect
  136. if the C{subCommands} attribute has been defined on the
  137. L{twisted.python.usage.Options} instance.
  138. """
  139. if genSubs and getattr(self.options, 'subCommands', None) is not None:
  140. gen = ZshArgumentsGenerator(self.options, self.cmdName, self.file)
  141. gen.extraActions.insert(0, SubcommandAction())
  142. gen.write()
  143. self.file.write(b'local _zsh_subcmds_array\n_zsh_subcmds_array=(\n')
  144. for (cmd, short, parser, desc) in self.options.subCommands:
  145. self.file.write(
  146. b'\"' + cmd.encode('utf-8') + b':' + desc.encode('utf-8') +b'\"\n')
  147. self.file.write(b")\n\n")
  148. self.file.write(b'_describe "sub-command" _zsh_subcmds_array\n')
  149. else:
  150. gen = ZshArgumentsGenerator(self.options, self.cmdName, self.file)
  151. gen.write()
  152. class ZshSubcommandBuilder(ZshBuilder):
  153. """
  154. Constructs zsh code that will complete options for a given usage.Options
  155. instance, and also for a single sub-command. This will only be used in
  156. the case where the user is completing options for a specific subcommand.
  157. @type subOptions: L{twisted.python.usage.Options}
  158. @ivar subOptions: The L{twisted.python.usage.Options} instance defined for
  159. the sub command.
  160. """
  161. def __init__(self, subOptions, *args):
  162. self.subOptions = subOptions
  163. ZshBuilder.__init__(self, *args)
  164. def write(self):
  165. """
  166. Generate the completion function and write it to the output file
  167. @return: L{None}
  168. """
  169. gen = ZshArgumentsGenerator(self.options, self.cmdName, self.file)
  170. gen.extraActions.insert(0, SubcommandAction())
  171. gen.write()
  172. gen = ZshArgumentsGenerator(self.subOptions, self.cmdName, self.file)
  173. gen.write()
  174. class ZshArgumentsGenerator(object):
  175. """
  176. Generate a call to the zsh _arguments completion function
  177. based on data in a usage.Options instance
  178. @type options: L{twisted.python.usage.Options}
  179. @ivar options: The L{twisted.python.usage.Options} instance to generate for
  180. @type cmdName: C{str}
  181. @ivar cmdName: The name of the command we're generating completions for.
  182. @type file: C{file}
  183. @ivar file: The C{file} to write the completion function to
  184. The following non-constructor variables are populated by this class
  185. with data gathered from the C{Options} instance passed in, and its
  186. base classes.
  187. @type descriptions: C{dict}
  188. @ivar descriptions: A dict mapping long option names to alternate
  189. descriptions. When this variable is defined, the descriptions
  190. contained here will override those descriptions provided in the
  191. optFlags and optParameters variables.
  192. @type multiUse: C{list}
  193. @ivar multiUse: An iterable containing those long option names which may
  194. appear on the command line more than once. By default, options will
  195. only be completed one time.
  196. @type mutuallyExclusive: C{list} of C{tuple}
  197. @ivar mutuallyExclusive: A sequence of sequences, with each sub-sequence
  198. containing those long option names that are mutually exclusive. That is,
  199. those options that cannot appear on the command line together.
  200. @type optActions: C{dict}
  201. @ivar optActions: A dict mapping long option names to shell "actions".
  202. These actions define what may be completed as the argument to the
  203. given option, and should be given as instances of
  204. L{twisted.python.usage.Completer}.
  205. Callables may instead be given for the values in this dict. The
  206. callable should accept no arguments, and return a C{Completer}
  207. instance used as the action.
  208. @type extraActions: C{list} of C{twisted.python.usage.Completer}
  209. @ivar extraActions: Extra arguments are those arguments typically
  210. appearing at the end of the command-line, which are not associated
  211. with any particular named option. That is, the arguments that are
  212. given to the parseArgs() method of your usage.Options subclass.
  213. """
  214. def __init__(self, options, cmdName, file):
  215. self.options = options
  216. self.cmdName = cmdName
  217. self.file = file
  218. self.descriptions = {}
  219. self.multiUse = set()
  220. self.mutuallyExclusive = []
  221. self.optActions = {}
  222. self.extraActions = []
  223. for cls in reversed(inspect.getmro(options.__class__)):
  224. data = getattr(cls, 'compData', None)
  225. if data:
  226. self.descriptions.update(data.descriptions)
  227. self.optActions.update(data.optActions)
  228. self.multiUse.update(data.multiUse)
  229. self.mutuallyExclusive.extend(data.mutuallyExclusive)
  230. # I don't see any sane way to aggregate extraActions, so just
  231. # take the one at the top of the MRO (nearest the `options'
  232. # instance).
  233. if data.extraActions:
  234. self.extraActions = data.extraActions
  235. aCL = reflect.accumulateClassList
  236. optFlags = []
  237. optParams = []
  238. aCL(options.__class__, 'optFlags', optFlags)
  239. aCL(options.__class__, 'optParameters', optParams)
  240. for i, optList in enumerate(optFlags):
  241. if len(optList) != 3:
  242. optFlags[i] = util.padTo(3, optList)
  243. for i, optList in enumerate(optParams):
  244. if len(optList) != 5:
  245. optParams[i] = util.padTo(5, optList)
  246. self.optFlags = optFlags
  247. self.optParams = optParams
  248. paramNameToDefinition = {}
  249. for optList in optParams:
  250. paramNameToDefinition[optList[0]] = optList[1:]
  251. self.paramNameToDefinition = paramNameToDefinition
  252. flagNameToDefinition = {}
  253. for optList in optFlags:
  254. flagNameToDefinition[optList[0]] = optList[1:]
  255. self.flagNameToDefinition = flagNameToDefinition
  256. allOptionsNameToDefinition = {}
  257. allOptionsNameToDefinition.update(paramNameToDefinition)
  258. allOptionsNameToDefinition.update(flagNameToDefinition)
  259. self.allOptionsNameToDefinition = allOptionsNameToDefinition
  260. self.addAdditionalOptions()
  261. # makes sure none of the Completions metadata references
  262. # option names that don't exist. (great for catching typos)
  263. self.verifyZshNames()
  264. self.excludes = self.makeExcludesDict()
  265. def write(self):
  266. """
  267. Write the zsh completion code to the file given to __init__
  268. @return: L{None}
  269. """
  270. self.writeHeader()
  271. self.writeExtras()
  272. self.writeOptions()
  273. self.writeFooter()
  274. def writeHeader(self):
  275. """
  276. This is the start of the code that calls _arguments
  277. @return: L{None}
  278. """
  279. self.file.write(b'#compdef ' + self.cmdName.encode('utf-8') +
  280. b'\n\n'
  281. b'_arguments -s -A "-*" \\\n')
  282. def writeOptions(self):
  283. """
  284. Write out zsh code for each option in this command
  285. @return: L{None}
  286. """
  287. optNames = list(self.allOptionsNameToDefinition.keys())
  288. optNames.sort()
  289. for longname in optNames:
  290. self.writeOpt(longname)
  291. def writeExtras(self):
  292. """
  293. Write out completion information for extra arguments appearing on the
  294. command-line. These are extra positional arguments not associated
  295. with a named option. That is, the stuff that gets passed to
  296. Options.parseArgs().
  297. @return: L{None}
  298. @raises: ValueError: if C{Completer} with C{repeat=True} is found and
  299. is not the last item in the C{extraActions} list.
  300. """
  301. for i, action in enumerate(self.extraActions):
  302. # a repeatable action must be the last action in the list
  303. if action._repeat and i != len(self.extraActions) - 1:
  304. raise ValueError("Completer with repeat=True must be "
  305. "last item in Options.extraActions")
  306. self.file.write(
  307. escape(action._shellCode('', usage._ZSH)).encode('utf-8'))
  308. self.file.write(b' \\\n')
  309. def writeFooter(self):
  310. """
  311. Write the last bit of code that finishes the call to _arguments
  312. @return: L{None}
  313. """
  314. self.file.write(b'&& return 0\n')
  315. def verifyZshNames(self):
  316. """
  317. Ensure that none of the option names given in the metadata are typoed
  318. @return: L{None}
  319. @raise ValueError: Raised if unknown option names have been found.
  320. """
  321. def err(name):
  322. raise ValueError("Unknown option name \"%s\" found while\n"
  323. "examining Completions instances on %s" % (
  324. name, self.options))
  325. for name in itertools.chain(self.descriptions, self.optActions,
  326. self.multiUse):
  327. if name not in self.allOptionsNameToDefinition:
  328. err(name)
  329. for seq in self.mutuallyExclusive:
  330. for name in seq:
  331. if name not in self.allOptionsNameToDefinition:
  332. err(name)
  333. def excludeStr(self, longname, buildShort=False):
  334. """
  335. Generate an "exclusion string" for the given option
  336. @type longname: C{str}
  337. @param longname: The long option name (e.g. "verbose" instead of "v")
  338. @type buildShort: C{bool}
  339. @param buildShort: May be True to indicate we're building an excludes
  340. string for the short option that corresponds to the given long opt.
  341. @return: The generated C{str}
  342. """
  343. if longname in self.excludes:
  344. exclusions = self.excludes[longname].copy()
  345. else:
  346. exclusions = set()
  347. # if longname isn't a multiUse option (can't appear on the cmd line more
  348. # than once), then we have to exclude the short option if we're
  349. # building for the long option, and vice versa.
  350. if longname not in self.multiUse:
  351. if buildShort is False:
  352. short = self.getShortOption(longname)
  353. if short is not None:
  354. exclusions.add(short)
  355. else:
  356. exclusions.add(longname)
  357. if not exclusions:
  358. return ''
  359. strings = []
  360. for optName in exclusions:
  361. if len(optName) == 1:
  362. # short option
  363. strings.append("-" + optName)
  364. else:
  365. strings.append("--" + optName)
  366. strings.sort() # need deterministic order for reliable unit-tests
  367. return "(%s)" % " ".join(strings)
  368. def makeExcludesDict(self):
  369. """
  370. @return: A C{dict} that maps each option name appearing in
  371. self.mutuallyExclusive to a list of those option names that is it
  372. mutually exclusive with (can't appear on the cmd line with).
  373. """
  374. #create a mapping of long option name -> single character name
  375. longToShort = {}
  376. for optList in itertools.chain(self.optParams, self.optFlags):
  377. if optList[1] != None:
  378. longToShort[optList[0]] = optList[1]
  379. excludes = {}
  380. for lst in self.mutuallyExclusive:
  381. for i, longname in enumerate(lst):
  382. tmp = set(lst[:i] + lst[i+1:])
  383. for name in tmp.copy():
  384. if name in longToShort:
  385. tmp.add(longToShort[name])
  386. if longname in excludes:
  387. excludes[longname] = excludes[longname].union(tmp)
  388. else:
  389. excludes[longname] = tmp
  390. return excludes
  391. def writeOpt(self, longname):
  392. """
  393. Write out the zsh code for the given argument. This is just part of the
  394. one big call to _arguments
  395. @type longname: C{str}
  396. @param longname: The long option name (e.g. "verbose" instead of "v")
  397. @return: L{None}
  398. """
  399. if longname in self.flagNameToDefinition:
  400. # It's a flag option. Not one that takes a parameter.
  401. longField = "--%s" % longname
  402. else:
  403. longField = "--%s=" % longname
  404. short = self.getShortOption(longname)
  405. if short != None:
  406. shortField = "-" + short
  407. else:
  408. shortField = ''
  409. descr = self.getDescription(longname)
  410. descriptionField = descr.replace("[", "\[")
  411. descriptionField = descriptionField.replace("]", "\]")
  412. descriptionField = '[%s]' % descriptionField
  413. actionField = self.getAction(longname)
  414. if longname in self.multiUse:
  415. multiField = '*'
  416. else:
  417. multiField = ''
  418. longExclusionsField = self.excludeStr(longname)
  419. if short:
  420. #we have to write an extra line for the short option if we have one
  421. shortExclusionsField = self.excludeStr(longname, buildShort=True)
  422. self.file.write(escape('%s%s%s%s%s' % (shortExclusionsField,
  423. multiField, shortField, descriptionField, actionField)).encode('utf-8'))
  424. self.file.write(b' \\\n')
  425. self.file.write(escape('%s%s%s%s%s' % (longExclusionsField,
  426. multiField, longField, descriptionField, actionField)).encode('utf-8'))
  427. self.file.write(b' \\\n')
  428. def getAction(self, longname):
  429. """
  430. Return a zsh "action" string for the given argument
  431. @return: C{str}
  432. """
  433. if longname in self.optActions:
  434. if callable(self.optActions[longname]):
  435. action = self.optActions[longname]()
  436. else:
  437. action = self.optActions[longname]
  438. return action._shellCode(longname, usage._ZSH)
  439. if longname in self.paramNameToDefinition:
  440. return ':%s:_files' % (longname,)
  441. return ''
  442. def getDescription(self, longname):
  443. """
  444. Return the description to be used for this argument
  445. @return: C{str}
  446. """
  447. #check if we have an alternate descr for this arg, and if so use it
  448. if longname in self.descriptions:
  449. return self.descriptions[longname]
  450. #otherwise we have to get it from the optFlags or optParams
  451. try:
  452. descr = self.flagNameToDefinition[longname][1]
  453. except KeyError:
  454. try:
  455. descr = self.paramNameToDefinition[longname][2]
  456. except KeyError:
  457. descr = None
  458. if descr is not None:
  459. return descr
  460. # let's try to get it from the opt_foo method doc string if there is one
  461. longMangled = longname.replace('-', '_') # this is what t.p.usage does
  462. obj = getattr(self.options, 'opt_%s' % longMangled, None)
  463. if obj is not None:
  464. descr = descrFromDoc(obj)
  465. if descr is not None:
  466. return descr
  467. return longname # we really ought to have a good description to use
  468. def getShortOption(self, longname):
  469. """
  470. Return the short option letter or None
  471. @return: C{str} or L{None}
  472. """
  473. optList = self.allOptionsNameToDefinition[longname]
  474. return optList[0] or None
  475. def addAdditionalOptions(self):
  476. """
  477. Add additional options to the optFlags and optParams lists.
  478. These will be defined by 'opt_foo' methods of the Options subclass
  479. @return: L{None}
  480. """
  481. methodsDict = {}
  482. reflect.accumulateMethods(self.options, methodsDict, 'opt_')
  483. methodToShort = {}
  484. for name in methodsDict.copy():
  485. if len(name) == 1:
  486. methodToShort[methodsDict[name]] = name
  487. del methodsDict[name]
  488. for methodName, methodObj in methodsDict.items():
  489. longname = methodName.replace('_', '-') # t.p.usage does this
  490. # if this option is already defined by the optFlags or
  491. # optParameters then we don't want to override that data
  492. if longname in self.allOptionsNameToDefinition:
  493. continue
  494. descr = self.getDescription(longname)
  495. short = None
  496. if methodObj in methodToShort:
  497. short = methodToShort[methodObj]
  498. reqArgs = methodObj.__func__.__code__.co_argcount
  499. if reqArgs == 2:
  500. self.optParams.append([longname, short, None, descr])
  501. self.paramNameToDefinition[longname] = [short, None, descr]
  502. self.allOptionsNameToDefinition[longname] = [short, None, descr]
  503. else:
  504. # reqArgs must equal 1. self.options would have failed
  505. # to instantiate if it had opt_ methods with bad signatures.
  506. self.optFlags.append([longname, short, descr])
  507. self.flagNameToDefinition[longname] = [short, descr]
  508. self.allOptionsNameToDefinition[longname] = [short, None, descr]
  509. def descrFromDoc(obj):
  510. """
  511. Generate an appropriate description from docstring of the given object
  512. """
  513. if obj.__doc__ is None or obj.__doc__.isspace():
  514. return None
  515. lines = [x.strip() for x in obj.__doc__.split("\n")
  516. if x and not x.isspace()]
  517. return " ".join(lines)
  518. def escape(x):
  519. """
  520. Shell escape the given string
  521. Implementation borrowed from now-deprecated commands.mkarg() in the stdlib
  522. """
  523. if '\'' not in x:
  524. return '\'' + x + '\''
  525. s = '"'
  526. for c in x:
  527. if c in '\\$"`':
  528. s = s + '\\'
  529. s = s + c
  530. s = s + '"'
  531. return s