Coverage for management/commands/autocomplete.py: 0%

134 statements  

« prev     ^ index     » next       coverage.py v7.3.4, created at 2024-01-20 17:58 +0000

1import sys 

2import os 

3from click.shell_completion import get_completion_class, CompletionItem, add_completion_class 

4from django_typer import TyperCommand, command, get_command 

5from django.core.management import CommandError, ManagementUtility 

6from django.utils.translation import gettext_lazy as _ 

7from django.utils.module_loading import import_string 

8from django.utils.functional import cached_property 

9import typing as t 

10from typer import Argument, Option, echo 

11from typer.completion import Shells, completion_init 

12from click.parser import split_arg_string 

13from pathlib import Path 

14import contextlib 

15import io 

16 

17 

18try: 

19 from shellingham import detect_shell 

20 detected_shell = detect_shell()[0] 

21except Exception: 

22 detected_shell = None 

23 

24 

25class Command(TyperCommand): 

26 """ 

27 This command installs autocompletion for the current shell. This command uses the typer/click 

28 autocompletion scripts to generate the autocompletion items, but monkey patches the scripts 

29 to invoke our bundled autocomplete script which fails over to the django autocomplete function 

30 when the command being completed is not a TyperCommand. When the django autocomplete function 

31 is used we also wrap it so that it works for any supported click/typer shell, not just bash. 

32 

33 We also provide a remove command to easily remove the installed script. 

34 

35 Great pains are taken to use the upstream dependency's autocomplete logic. This is so advances 

36 and additional shell support implemented upstream should just work. However, it would be possible 

37 to add support for new shells here using the pluggable logic that click provides. It is  

38 probably a better idea however to add support for new shells at the typer level. 

39 

40 Shell autocompletion can be brittle with every shell having its own quirks and nuances. We 

41 make a good faith effort here to support all the shells that typer/click support, but there 

42 can easily be system specific configuration issues that prevent this from working. In those 

43 cases users should refer to the online documentation for their specific shell to troubleshoot. 

44 """ 

45 

46 help = _('Install autocompletion for the current shell.') 

47 

48 """ 

49 The name of the django manage command script to install autocompletion for. We do 

50 not want to hardcode this as 'manage.py' because users are free to rename and implement 

51 their own manage scripts! The safest way to do this is therefore to require this install 

52 script to be a management command itself and then fetch the name of the script that invoked 

53 it. 

54 """ 

55 MANAGE_SCRIPT = Path(sys.argv[0]).name 

56 

57 """ 

58 The name of the script that will generate the autocompletion items from the tabbed command string. 

59 By default we use our bundled script, but allow users to override this. 

60 """ 

61 COMPLETION_SCRIPT = f'{MANAGE_SCRIPT} {Path(__file__).name} complete' 

62 

63 # disable the system checks - no reason to run these for this one-off command 

64 requires_system_checks = [] 

65 requires_migrations_checks = False 

66 

67 # remove unnecessary django command base parameters - these just clutter the help 

68 django_params = [ 

69 cmd for cmd in 

70 TyperCommand.django_params 

71 if cmd not in [ 

72 'version', 'skip_checks', 

73 'no_color', 'force_color' 

74 ] 

75 ] 

76 

77 _shell: Shells 

78 

79 COMPLETE_VAR = '_DJANGO_COMPLETE' 

80 

81 @property 

82 def shell(self): 

83 """ 

84 Get the active shell. If not explicitly set, it first tries to find the shell 

85 in the environment variable autocomplete scripts set and failing that it will try 

86 to autodetect the shell. 

87 """ 

88 return getattr( 

89 self, 

90 '_shell', 

91 Shells( 

92 os.environ[self.COMPLETE_VAR].partition('_')[2] 

93 if self.COMPLETE_VAR in os.environ else 

94 detect_shell()[0] 

95 ) 

96 ) 

97 

98 @shell.setter 

99 def shell(self, shell): 

100 """Set the shell to install autocompletion for.""" 

101 if shell is None: 

102 raise CommandError(_( 

103 'Unable to detect shell. Please specify the shell to install or remove ' 

104 'autocompletion for.' 

105 )) 

106 self._shell = shell if isinstance(shell, Shells) else Shells(shell) 

107 

108 @cached_property 

109 def noop_command(self): 

110 """ 

111 This is a no-op command that is used to bootstrap click Completion classes. It 

112 has no use other than to avoid any potential attribute errors when we emulate 

113 upstream completion logic 

114 """ 

115 return self.get_subcommand('noop') 

116 

117 def patch_script( 

118 self, 

119 shell: t.Optional[Shells] = None, 

120 command: str=COMPLETION_SCRIPT, 

121 fallback: t.Optional[str] = None 

122 ) -> None: 

123 """ 

124 We have to monkey patch the typer completion scripts to point to our custom 

125 autocomplete script. This is potentially brittle but thats why we have robust 

126 CI! 

127 

128 :param shell: The shell to patch the completion script for. 

129 :param command: The name of the autocomplete script to use - defaults to the bundled script. 

130 :param fallback: The python import path to a fallback autocomplete function to use when 

131 the completion command is not a TyperCommand. Defaults to None, which means the bundled 

132 complete script will fallback to the django autocomplete function, but wrap it so it works 

133 for all supported shells other than just bash. 

134 """ 

135 # do not import this private stuff until we need it - avoids it tanking the whole 

136 # library if these imports change 

137 from typer import _completion_shared as typer_scripts 

138 shell = shell or self.shell 

139 

140 fallback = f' --fallback {fallback}' if fallback else '' 

141 complete_cmd = f'{command}{fallback}' 

142 

143 def replace(s: str, old: str, new: str, occurrences: list[int]) -> str: 

144 """ 

145 :param s: The string to modify 

146 :param old: The string to replace 

147 :param new: The string to replace with 

148 :param occurrences: A list of occurrences of the old string to replace with the 

149 new string, where the occurrence number is the zero-based count of the old 

150 strings appearance in the string when counting from the front. 

151 """ 

152 count = 0 

153 result = "" 

154 start = 0 

155 

156 for end in range(len(s)): 

157 if s[start:end+1].endswith(old): 

158 if count in occurrences: 

159 result += f'{s[start:end+1-len(old)]}{new}' 

160 start = end + 1 

161 else: 

162 result += s[start:end+1] 

163 start = end + 1 

164 count += 1 

165 

166 result += s[start:] 

167 return result 

168 

169 if shell is Shells.bash: 

170 typer_scripts._completion_scripts[Shells.bash.value] = replace( 

171 typer_scripts.COMPLETION_SCRIPT_BASH, 

172 '%(prog_name)s', 

173 f'$0 autocomplete complete', 

174 [0] 

175 ) 

176 elif shell is Shells.zsh: 

177 typer_scripts._completion_scripts[Shells.zsh.value] = replace( 

178 typer_scripts.COMPLETION_SCRIPT_ZSH, 

179 '%(prog_name)s', 

180 f'${ words[0,1]} autocomplete complete', 

181 [1] 

182 ) 

183 # TODO: Figure out how to patch these 

184 # elif shell is Shells.fish: 

185 # typer_scripts.COMPLETION_SCRIPT_FISH = replace_last( 

186 # typer_scripts.COMPLETION_SCRIPT_FISH, 

187 # '%(prog_name)s', 

188 # 'django_complete' 

189 # ) 

190 # elif shell is Shells.powershell: 

191 # typer_scripts.COMPLETION_SCRIPT_POWER_SHELL = replace_last( 

192 # typer_scripts.COMPLETION_SCRIPT_POWER_SHELL, 

193 # '%(prog_name)s', 

194 # 'django_complete' 

195 # ) 

196 

197 @command(help=_('Install autocompletion.')) 

198 def install( 

199 self, 

200 shell: t.Annotated[ 

201 t.Optional[Shells], 

202 Argument(help=_('Specify the shell to install or remove autocompletion for.')) 

203 ] = detected_shell, 

204 completion_script: t.Annotated[ 

205 str, 

206 Option(help=_('The name of the autocomplete script.') 

207 )] = COMPLETION_SCRIPT, 

208 manage_script: t.Annotated[ 

209 str, 

210 Option(help=_('The name of the django manage script to install autocompletion for.') 

211 )] = MANAGE_SCRIPT, 

212 fallback: t.Annotated[ 

213 t.Optional[str], 

214 Option(help=_( 

215 'The python import path to a fallback autocomplete function to use when ' 

216 'the completion command is not a TyperCommand.' 

217 ) 

218 )] = None 

219 ): 

220 # do not import this private stuff until we need it - avoids tanking the whole 

221 # library if these imports change 

222 from typer._completion_shared import install 

223 self.shell = shell 

224 self.patch_script(command=completion_script, fallback=fallback) 

225 install_path = install( 

226 shell=self.shell.value, 

227 prog_name=manage_script, 

228 complete_var=self.COMPLETE_VAR 

229 )[1] 

230 self.stdout.write(self.style.SUCCESS( 

231 _('Installed autocompletion for %(shell)s @ %(install_path)s') % { 

232 'shell': shell.value, 'install_path': install_path 

233 } 

234 )) 

235 

236 @command(help=_('Remove autocompletion.')) 

237 def remove( 

238 self, 

239 shell: t.Annotated[ 

240 t.Optional[Shells], 

241 Argument(help=_('Specify the shell to install or remove autocompletion for.')) 

242 ] = detected_shell 

243 ): 

244 from typer._completion_shared import install 

245 # its less brittle to install and use the returned path to uninstall 

246 self.shell = shell 

247 stdout = self.stdout 

248 self.stdout = io.StringIO() 

249 installed_path = install(shell=self.shell.value, prog_name=self.COMPLETION_SCRIPT)[1] 

250 installed_path.unlink() 

251 self.stdout = stdout 

252 self.stdout.write(self.style.WARNING( 

253 _('Removed autocompletion for %(shell)s.') % {'shell': shell.value} 

254 )) 

255 

256 @command(help=_('Generate autocompletion for command string.'), hidden=False) 

257 def complete( 

258 self, 

259 command: t.Annotated[ 

260 t.Optional[str], 

261 Argument( 

262 help=_('The command string to generate completion suggestions for.') 

263 ) 

264 ] = None, 

265 fallback: t.Annotated[ 

266 t.Optional[str], 

267 Option( 

268 help=_( 

269 'The python import path to a fallback autocomplete function to use when ' 

270 'the completion command is not a TyperCommand. By default, the builtin ' 

271 'django autocomplete function is used.' 

272 ) 

273 ) 

274 ] = None 

275 ): 

276 """ 

277 We implement the autocomplete generation script as a Django command because the 

278 Django environment needs to be bootstrapped for it to work. This also allows 

279 us to test autocompletions in a platform agnostic way. 

280 """ 

281 completion_init() 

282 CompletionClass = get_completion_class(self.shell.value) 

283 if command: 

284 # when the command is given, this is a user testing their autocompletion, 

285 # so we need to override the completion classes get_completion_args logic 

286 # because our entry point was not an installed completion script 

287 def get_completion_args(self) -> t.Tuple[t.List[str], str]: 

288 cwords = split_arg_string(command) 

289 # allow users to not specify the manage script, but allow for it 

290 # if they do by lopping it off - same behavior as upstream classes 

291 try: 

292 if Path(cwords[0]).resolve() == Path(sys.argv[0]).resolve(): 

293 cwords = cwords[1:] 

294 except Exception: 

295 pass 

296 return cwords, cwords[-1] if len(cwords) and not command[-1].isspace() else '' 

297 

298 CompletionClass.get_completion_args = get_completion_args 

299 add_completion_class(self.shell.value, CompletionClass) 

300 

301 args, incomplete = CompletionClass( 

302 cli=self.noop_command, 

303 ctx_args=self.noop_command, 

304 prog_name=self.MANAGE_SCRIPT, 

305 complete_var=self.COMPLETE_VAR 

306 ).get_completion_args() 

307 

308 with open('args.txt', 'w') as f: 

309 f.write(str(args)) 

310 f.write(incomplete) 

311 

312 fallback = import_string(fallback) if fallback else self.django_fallback 

313 

314 if not args: 

315 fallback() 

316 else: 

317 try: 

318 cmd = get_command(args[0]) 

319 if isinstance(cmd, TyperCommand): 

320 # invoking the command will trigger the autocompletion? 

321 #print('RUN') 

322 cmd() 

323 else: 

324 print('ELSE') 

325 fallback() 

326 except Exception as e: 

327 fallback() 

328 

329 def django_fallback(self): 

330 """ 

331 Run django's builtin bash autocomplete function. We wrap the click 

332 completion class to make it work for all supported shells, not just 

333 bash. 

334 """ 

335 CompletionClass = get_completion_class(self.shell.value) 

336 def get_completions(self, args, incomplete): 

337 # spoof bash environment variables 

338 # the first one is lopped off, so we insert a placeholder 0 

339 args = ['0', *args] 

340 if args[-1] != incomplete: 

341 args.append(incomplete) 

342 os.environ['COMP_WORDS'] = ' '.join(args) 

343 os.environ['COMP_CWORD'] = str(args.index(incomplete) + 1) 

344 os.environ['DJANGO_AUTO_COMPLETE'] = '1' 

345 dj_manager = ManagementUtility(args) 

346 capture_completions = io.StringIO() 

347 try: 

348 with contextlib.redirect_stdout(capture_completions): 

349 dj_manager.autocomplete() 

350 except SystemExit: 

351 completions = [ 

352 CompletionItem(item) 

353 for item in capture_completions.getvalue().split('\n') 

354 if item 

355 ] 

356 return completions 

357 CompletionClass.get_completions = get_completions 

358 echo( 

359 CompletionClass( 

360 cli=self.noop_command, 

361 ctx_args={}, 

362 prog_name=self.MANAGE_SCRIPT, 

363 complete_var=self.COMPLETE_VAR 

364 ).complete() 

365 ) 

366 

367 @command( 

368 hidden=True, 

369 context_settings={ 

370 'ignore_unknown_options': True, 

371 'allow_extra_args': True, 

372 'allow_interspersed_args': True 

373 } 

374 ) 

375 def noop(self): 

376 """ 

377 This is a no-op command that is used to bootstrap click Completion classes. It 

378 has no use other than to avoid any potential attribute errors when we spoof 

379 completion logic 

380 """ 

381 pass