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
« 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
18try:
19 from shellingham import detect_shell
20 detected_shell = detect_shell()[0]
21except Exception:
22 detected_shell = None
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.
33 We also provide a remove command to easily remove the installed script.
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.
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 """
46 help = _('Install autocompletion for the current shell.')
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
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'
63 # disable the system checks - no reason to run these for this one-off command
64 requires_system_checks = []
65 requires_migrations_checks = False
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 ]
77 _shell: Shells
79 COMPLETE_VAR = '_DJANGO_COMPLETE'
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 )
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)
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')
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!
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
140 fallback = f' --fallback {fallback}' if fallback else ''
141 complete_cmd = f'{command}{fallback}'
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
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
166 result += s[start:]
167 return result
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 # )
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 ))
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 ))
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 ''
298 CompletionClass.get_completion_args = get_completion_args
299 add_completion_class(self.shell.value, CompletionClass)
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()
308 with open('args.txt', 'w') as f:
309 f.write(str(args))
310 f.write(incomplete)
312 fallback = import_string(fallback) if fallback else self.django_fallback
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()
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 )
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