Coverage for typer/_completion_shared.py: 100%

105 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-09-09 18:26 +0000

1import os 1habfgcde

2import re 1habfgcde

3import subprocess 1habfgcde

4from enum import Enum 1habfgcde

5from pathlib import Path 1habfgcde

6from typing import Optional, Tuple 1habfgcde

7 

8import click 1habfgcde

9 

10try: 1habfgcde

11 import shellingham 1habfgcde

12except ImportError: # pragma: no cover 

13 shellingham = None 

14 

15 

16class Shells(str, Enum): 1habfgcde

17 bash = "bash" 1habfgcde

18 zsh = "zsh" 1habfgcde

19 fish = "fish" 1habfgcde

20 powershell = "powershell" 1habfgcde

21 pwsh = "pwsh" 1habfgcde

22 

23 

24COMPLETION_SCRIPT_BASH = """ 1abfgcde

25%(complete_func)s() { 

26 local IFS=$'\n' 

27 COMPREPLY=( $( env COMP_WORDS="${COMP_WORDS[*]}" \\ 

28 COMP_CWORD=$COMP_CWORD \\ 

29 %(autocomplete_var)s=complete_bash $1 ) ) 

30 return 0 

31} 

32 

33complete -o default -F %(complete_func)s %(prog_name)s 

34""" 

35 

36COMPLETION_SCRIPT_ZSH = """ 1abfgcde

37#compdef %(prog_name)s 

38 

39%(complete_func)s() { 

40 eval $(env _TYPER_COMPLETE_ARGS="${words[1,$CURRENT]}" %(autocomplete_var)s=complete_zsh %(prog_name)s) 

41} 

42 

43compdef %(complete_func)s %(prog_name)s 

44""" 

45 

46COMPLETION_SCRIPT_FISH = 'complete --command %(prog_name)s --no-files --arguments "(env %(autocomplete_var)s=complete_fish _TYPER_COMPLETE_FISH_ACTION=get-args _TYPER_COMPLETE_ARGS=(commandline -cp) %(prog_name)s)" --condition "env %(autocomplete_var)s=complete_fish _TYPER_COMPLETE_FISH_ACTION=is-args _TYPER_COMPLETE_ARGS=(commandline -cp) %(prog_name)s"' 1habfgcde

47 

48COMPLETION_SCRIPT_POWER_SHELL = """ 1abfgcde

49Import-Module PSReadLine 

50Set-PSReadLineKeyHandler -Chord Tab -Function MenuComplete 

51$scriptblock = { 

52 param($wordToComplete, $commandAst, $cursorPosition) 

53 $Env:%(autocomplete_var)s = "complete_powershell" 

54 $Env:_TYPER_COMPLETE_ARGS = $commandAst.ToString() 

55 $Env:_TYPER_COMPLETE_WORD_TO_COMPLETE = $wordToComplete 

56 %(prog_name)s | ForEach-Object { 

57 $commandArray = $_ -Split ":::" 

58 $command = $commandArray[0] 

59 $helpString = $commandArray[1] 

60 [System.Management.Automation.CompletionResult]::new( 

61 $command, $command, 'ParameterValue', $helpString) 

62 } 

63 $Env:%(autocomplete_var)s = "" 

64 $Env:_TYPER_COMPLETE_ARGS = "" 

65 $Env:_TYPER_COMPLETE_WORD_TO_COMPLETE = "" 

66} 

67Register-ArgumentCompleter -Native -CommandName %(prog_name)s -ScriptBlock $scriptblock 

68""" 

69 

70_completion_scripts = { 1abfgcde

71 "bash": COMPLETION_SCRIPT_BASH, 

72 "zsh": COMPLETION_SCRIPT_ZSH, 

73 "fish": COMPLETION_SCRIPT_FISH, 

74 "powershell": COMPLETION_SCRIPT_POWER_SHELL, 

75 "pwsh": COMPLETION_SCRIPT_POWER_SHELL, 

76} 

77 

78# TODO: Probably refactor this, copied from Click 7.x 

79_invalid_ident_char_re = re.compile(r"[^a-zA-Z0-9_]") 1habfgcde

80 

81 

82def get_completion_script(*, prog_name: str, complete_var: str, shell: str) -> str: 1habfgcde

83 cf_name = _invalid_ident_char_re.sub("", prog_name.replace("-", "_")) 1habfgcde

84 script = _completion_scripts.get(shell) 1habfgcde

85 if script is None: 1habfgcde

86 click.echo(f"Shell {shell} not supported.", err=True) 1habfgcde

87 raise click.exceptions.Exit(1) 1habfgcde

88 return ( 1abfgcde

89 script 

90 % { 

91 "complete_func": f"_{cf_name}_completion", 

92 "prog_name": prog_name, 

93 "autocomplete_var": complete_var, 

94 } 

95 ).strip() 

96 

97 

98def install_bash(*, prog_name: str, complete_var: str, shell: str) -> Path: 1habfgcde

99 # Ref: https://github.com/scop/bash-completion#faq 

100 # It seems bash-completion is the official completion system for bash: 

101 # Ref: https://www.gnu.org/software/bash/manual/html_node/A-Programmable-Completion-Example.html 

102 # But installing in the locations from the docs doesn't seem to have effect 

103 completion_path = Path.home() / ".bash_completions" / f"{prog_name}.sh" 1habfgcde

104 rc_path = Path.home() / ".bashrc" 1habfgcde

105 rc_path.parent.mkdir(parents=True, exist_ok=True) 1habfgcde

106 rc_content = "" 1habfgcde

107 if rc_path.is_file(): 1habfgcde

108 rc_content = rc_path.read_text() 1habcde

109 completion_init_lines = [f"source '{completion_path}'"] 1habfgcde

110 for line in completion_init_lines: 1habfgcde

111 if line not in rc_content: # pragma: no cover 1habfgcde

112 rc_content += f"\n{line}" 1habfgcde

113 rc_content += "\n" 1habfgcde

114 rc_path.write_text(rc_content) 1habfgcde

115 # Install completion 

116 completion_path.parent.mkdir(parents=True, exist_ok=True) 1habfgcde

117 script_content = get_completion_script( 1habfgcde

118 prog_name=prog_name, complete_var=complete_var, shell=shell 

119 ) 

120 completion_path.write_text(script_content) 1habfgcde

121 return completion_path 1habfgcde

122 

123 

124def install_zsh(*, prog_name: str, complete_var: str, shell: str) -> Path: 1habfgcde

125 # Setup Zsh and load ~/.zfunc 

126 zshrc_path = Path.home() / ".zshrc" 1habfgcde

127 zshrc_path.parent.mkdir(parents=True, exist_ok=True) 1habfgcde

128 zshrc_content = "" 1habfgcde

129 if zshrc_path.is_file(): 1habfgcde

130 zshrc_content = zshrc_path.read_text() 1habfgcde

131 completion_line = "fpath+=~/.zfunc; autoload -Uz compinit; compinit" 1habfgcde

132 if completion_line not in zshrc_content: 1habfgcde

133 zshrc_content += f"\n{completion_line}\n" 1habfgcde

134 style_line = "zstyle ':completion:*' menu select" 1habfgcde

135 # TODO: consider setting the style only for the current program 

136 # style_line = f"zstyle ':completion:*:*:{prog_name}:*' menu select" 

137 # Install zstyle completion config only if the user doesn't have a customization 

138 if "zstyle" not in zshrc_content: 1habfgcde

139 zshrc_content += f"\n{style_line}\n" 1habfgcde

140 zshrc_content = f"{zshrc_content.strip()}\n" 1habfgcde

141 zshrc_path.write_text(zshrc_content) 1habfgcde

142 # Install completion under ~/.zfunc/ 

143 path_obj = Path.home() / f".zfunc/_{prog_name}" 1habfgcde

144 path_obj.parent.mkdir(parents=True, exist_ok=True) 1habfgcde

145 script_content = get_completion_script( 1habfgcde

146 prog_name=prog_name, complete_var=complete_var, shell=shell 

147 ) 

148 path_obj.write_text(script_content) 1habfgcde

149 return path_obj 1habfgcde

150 

151 

152def install_fish(*, prog_name: str, complete_var: str, shell: str) -> Path: 1habfgcde

153 path_obj = Path.home() / f".config/fish/completions/{prog_name}.fish" 1habfgcde

154 parent_dir: Path = path_obj.parent 1habfgcde

155 parent_dir.mkdir(parents=True, exist_ok=True) 1habfgcde

156 script_content = get_completion_script( 1habfgcde

157 prog_name=prog_name, complete_var=complete_var, shell=shell 

158 ) 

159 path_obj.write_text(f"{script_content}\n") 1habfgcde

160 return path_obj 1habfgcde

161 

162 

163def install_powershell(*, prog_name: str, complete_var: str, shell: str) -> Path: 1habfgcde

164 subprocess.run( 1habfgcde

165 [ 

166 shell, 

167 "-Command", 

168 "Set-ExecutionPolicy", 

169 "Unrestricted", 

170 "-Scope", 

171 "CurrentUser", 

172 ] 

173 ) 

174 result = subprocess.run( 1habfgcde

175 [shell, "-NoProfile", "-Command", "echo", "$profile"], 

176 check=True, 

177 stdout=subprocess.PIPE, 

178 ) 

179 if result.returncode != 0: # pragma: no cover 1habfgcde

180 click.echo("Couldn't get PowerShell user profile", err=True) 

181 raise click.exceptions.Exit(result.returncode) 

182 path_str = "" 1habfgcde

183 if isinstance(result.stdout, str): # pragma: no cover 1habfgcde

184 path_str = result.stdout 

185 if isinstance(result.stdout, bytes): 1habfgcde

186 try: 1habfgcde

187 # PowerShell would be predominant in Windows 

188 path_str = result.stdout.decode("windows-1252") 1habfgcde

189 except UnicodeDecodeError: # pragma: no cover 

190 try: 

191 path_str = result.stdout.decode("utf8") 

192 except UnicodeDecodeError: 

193 click.echo("Couldn't decode the path automatically", err=True) 

194 raise 

195 path_obj = Path(path_str.strip()) 1habfgcde

196 parent_dir: Path = path_obj.parent 1habfgcde

197 parent_dir.mkdir(parents=True, exist_ok=True) 1habfgcde

198 script_content = get_completion_script( 1habfgcde

199 prog_name=prog_name, complete_var=complete_var, shell=shell 

200 ) 

201 with path_obj.open(mode="a") as f: 1habfgcde

202 f.write(f"{script_content}\n") 1habfgcde

203 return path_obj 1habfgcde

204 

205 

206def install( 1abfgcde

207 shell: Optional[str] = None, 

208 prog_name: Optional[str] = None, 

209 complete_var: Optional[str] = None, 

210) -> Tuple[str, Path]: 

211 prog_name = prog_name or click.get_current_context().find_root().info_name 1habfgcde

212 assert prog_name 1habfgcde

213 if complete_var is None: 1habfgcde

214 complete_var = "_{}_COMPLETE".format(prog_name.replace("-", "_").upper()) 1habfgcde

215 test_disable_detection = os.getenv("_TYPER_COMPLETE_TEST_DISABLE_SHELL_DETECTION") 1habfgcde

216 if shell is None and shellingham is not None and not test_disable_detection: 1habfgcde

217 shell, _ = shellingham.detect_shell() 1habfgcde

218 if shell == "bash": 1habfgcde

219 installed_path = install_bash( 1habfgcde

220 prog_name=prog_name, complete_var=complete_var, shell=shell 

221 ) 

222 return shell, installed_path 1habfgcde

223 elif shell == "zsh": 1habfgcde

224 installed_path = install_zsh( 1habfgcde

225 prog_name=prog_name, complete_var=complete_var, shell=shell 

226 ) 

227 return shell, installed_path 1habfgcde

228 elif shell == "fish": 1habfgcde

229 installed_path = install_fish( 1habfgcde

230 prog_name=prog_name, complete_var=complete_var, shell=shell 

231 ) 

232 return shell, installed_path 1habfgcde

233 elif shell in {"powershell", "pwsh"}: 1habfgcde

234 installed_path = install_powershell( 1habfgcde

235 prog_name=prog_name, complete_var=complete_var, shell=shell 

236 ) 

237 return shell, installed_path 1habfgcde

238 else: 

239 click.echo(f"Shell {shell} is not supported.") 1habfgcde

240 raise click.exceptions.Exit(1) 1habfgcde