twisted/python/_shellcomp.py

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