Add tab completion for %model

Author: asmeurer

mypython/completion.py

@@ -35,7 +35,7 @@
 import traceback
 
 from .dircompletion import DirCompleter
-from .magic import MAGICS
+from .magic import MAGICS, MAGIC_COMPLETIONS
 
 def get_jedi_script_from_document(document, _locals, _globals, session):
     import jedi  # We keep this import in-line, to improve start-up time.
@@ -99,6 +99,11 @@ def get_completions(self, document, complete_event):
         for magic in MAGICS:
             if text_before_cursor.startswith(magic + ' '):
                 text_before_cursor = text_before_cursor[len(magic) + 1:]
+                if magic in MAGIC_COMPLETIONS:
+                    for completion in MAGIC_COMPLETIONS[magic]():
+                        yield Completion(completion,
+                                         -len(document.text_before_cursor), display_meta=magic)
+                    return
                 break
 
         if complete_event.completion_requested or self._complete_python_while_typing(document):

mypython/magic.py

@@ -58,6 +58,21 @@ def inner(rest):
     NON_PYTHON_MAGICS.append('%' + f.__name__[:-len('_magic')])
     return inner
 
+MAGIC_COMPLETIONS = {}
+
+def completions(get_completions):
+    """
+    Use like
+
+    @completions(lambda: ['a', 'b', 'c'])
+    def x_magic(rest):
+        ...
+    """
+    def wrapper(f):
+        MAGIC_COMPLETIONS['%' + f.__name__[:-len('_magic')]] = get_completions
+        return f
+    return wrapper
+
 def error(message):
     return """\
 import sys as _sys
@@ -402,14 +417,20 @@ def timings_magic(rest):
         print(f'{{_i:2d}} {{_format_time(TIMINGS[_i])}}')
 """
 
+def get_ai_models():
+    yield from list(ai.MODELS)
+    for model in ai.MODELS:
+        yield from ai.MODELS[model]['model_aliases']
+
+@completions(get_ai_models)
 @nonpython
 def model_magic(rest):
     """
     Set the current model for the completion engine.
     """
     rest = rest.strip()
     if not rest:
-        return error("no model specified")
+        return error("no model specified: available models are %s" % ', '.join(ai.MODELS))
 
     for model in ai.MODELS:
         if rest == model: