autocompletion.py 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176
  1. """Logic that powers autocompletion installed by ``pip completion``.
  2. """
  3. import optparse
  4. import os
  5. import sys
  6. from itertools import chain
  7. from typing import Any, Iterable, List, Optional
  8. from pip._internal.cli.main_parser import create_main_parser
  9. from pip._internal.commands import commands_dict, create_command
  10. from pip._internal.metadata import get_default_environment
  11. def autocomplete() -> None:
  12. """Entry Point for completion of main and subcommand options."""
  13. # Don't complete if user hasn't sourced bash_completion file.
  14. if "PIP_AUTO_COMPLETE" not in os.environ:
  15. return
  16. # Don't complete if autocompletion environment variables
  17. # are not present
  18. if not os.environ.get("COMP_WORDS") or not os.environ.get("COMP_CWORD"):
  19. return
  20. cwords = os.environ["COMP_WORDS"].split()[1:]
  21. cword = int(os.environ["COMP_CWORD"])
  22. try:
  23. current = cwords[cword - 1]
  24. except IndexError:
  25. current = ""
  26. parser = create_main_parser()
  27. subcommands = list(commands_dict)
  28. options = []
  29. # subcommand
  30. subcommand_name: Optional[str] = None
  31. for word in cwords:
  32. if word in subcommands:
  33. subcommand_name = word
  34. break
  35. # subcommand options
  36. if subcommand_name is not None:
  37. # special case: 'help' subcommand has no options
  38. if subcommand_name == "help":
  39. sys.exit(1)
  40. # special case: list locally installed dists for show and uninstall
  41. should_list_installed = not current.startswith("-") and subcommand_name in [
  42. "show",
  43. "uninstall",
  44. ]
  45. if should_list_installed:
  46. env = get_default_environment()
  47. lc = current.lower()
  48. installed = [
  49. dist.canonical_name
  50. for dist in env.iter_installed_distributions(local_only=True)
  51. if dist.canonical_name.startswith(lc)
  52. and dist.canonical_name not in cwords[1:]
  53. ]
  54. # if there are no dists installed, fall back to option completion
  55. if installed:
  56. for dist in installed:
  57. print(dist)
  58. sys.exit(1)
  59. should_list_installables = (
  60. not current.startswith("-") and subcommand_name == "install"
  61. )
  62. if should_list_installables:
  63. for path in auto_complete_paths(current, "path"):
  64. print(path)
  65. sys.exit(1)
  66. subcommand = create_command(subcommand_name)
  67. for opt in subcommand.parser.option_list_all:
  68. if opt.help != optparse.SUPPRESS_HELP:
  69. options += [
  70. (opt_str, opt.nargs) for opt_str in opt._long_opts + opt._short_opts
  71. ]
  72. # filter out previously specified options from available options
  73. prev_opts = [x.split("=")[0] for x in cwords[1 : cword - 1]]
  74. options = [(x, v) for (x, v) in options if x not in prev_opts]
  75. # filter options by current input
  76. options = [(k, v) for k, v in options if k.startswith(current)]
  77. # get completion type given cwords and available subcommand options
  78. completion_type = get_path_completion_type(
  79. cwords,
  80. cword,
  81. subcommand.parser.option_list_all,
  82. )
  83. # get completion files and directories if ``completion_type`` is
  84. # ``<file>``, ``<dir>`` or ``<path>``
  85. if completion_type:
  86. paths = auto_complete_paths(current, completion_type)
  87. options = [(path, 0) for path in paths]
  88. for option in options:
  89. opt_label = option[0]
  90. # append '=' to options which require args
  91. if option[1] and option[0][:2] == "--":
  92. opt_label += "="
  93. print(opt_label)
  94. else:
  95. # show main parser options only when necessary
  96. opts = [i.option_list for i in parser.option_groups]
  97. opts.append(parser.option_list)
  98. flattened_opts = chain.from_iterable(opts)
  99. if current.startswith("-"):
  100. for opt in flattened_opts:
  101. if opt.help != optparse.SUPPRESS_HELP:
  102. subcommands += opt._long_opts + opt._short_opts
  103. else:
  104. # get completion type given cwords and all available options
  105. completion_type = get_path_completion_type(cwords, cword, flattened_opts)
  106. if completion_type:
  107. subcommands = list(auto_complete_paths(current, completion_type))
  108. print(" ".join([x for x in subcommands if x.startswith(current)]))
  109. sys.exit(1)
  110. def get_path_completion_type(
  111. cwords: List[str], cword: int, opts: Iterable[Any]
  112. ) -> Optional[str]:
  113. """Get the type of path completion (``file``, ``dir``, ``path`` or None)
  114. :param cwords: same as the environmental variable ``COMP_WORDS``
  115. :param cword: same as the environmental variable ``COMP_CWORD``
  116. :param opts: The available options to check
  117. :return: path completion type (``file``, ``dir``, ``path`` or None)
  118. """
  119. if cword < 2 or not cwords[cword - 2].startswith("-"):
  120. return None
  121. for opt in opts:
  122. if opt.help == optparse.SUPPRESS_HELP:
  123. continue
  124. for o in str(opt).split("/"):
  125. if cwords[cword - 2].split("=")[0] == o:
  126. if not opt.metavar or any(
  127. x in ("path", "file", "dir") for x in opt.metavar.split("/")
  128. ):
  129. return opt.metavar
  130. return None
  131. def auto_complete_paths(current: str, completion_type: str) -> Iterable[str]:
  132. """If ``completion_type`` is ``file`` or ``path``, list all regular files
  133. and directories starting with ``current``; otherwise only list directories
  134. starting with ``current``.
  135. :param current: The word to be completed
  136. :param completion_type: path completion type(``file``, ``path`` or ``dir``)
  137. :return: A generator of regular files and/or directories
  138. """
  139. directory, filename = os.path.split(current)
  140. current_path = os.path.abspath(directory)
  141. # Don't complete paths if they can't be accessed
  142. if not os.access(current_path, os.R_OK):
  143. return
  144. filename = os.path.normcase(filename)
  145. # list all files that start with ``filename``
  146. file_list = (
  147. x for x in os.listdir(current_path) if os.path.normcase(x).startswith(filename)
  148. )
  149. for f in file_list:
  150. opt = os.path.join(current_path, f)
  151. comp_file = os.path.normcase(os.path.join(directory, f))
  152. # complete regular files when there is not ``<dir>`` after option
  153. # complete directories when there is ``<file>``, ``<path>`` or
  154. # ``<dir>``after option
  155. if completion_type != "dir" and os.path.isfile(opt):
  156. yield comp_file
  157. elif os.path.isdir(opt):
  158. yield os.path.join(comp_file, "")