Skip to content

Commit

Permalink
Add tab completion for %model
Browse files Browse the repository at this point in the history
  • Loading branch information
asmeurer committed Sep 12, 2024
1 parent d18ff5c commit 5394055
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 2 deletions.
7 changes: 6 additions & 1 deletion mypython/completion.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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):
Expand Down
23 changes: 22 additions & 1 deletion mypython/magic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 5394055

Please sign in to comment.