WIP: improve selection of environments

Author: blueyed

autoload/jedi.vim

@@ -34,6 +34,8 @@ let s:default_settings = {
     \ 'popup_select_first': 1,
     \ 'quickfix_window_height': 10,
     \ 'force_py_version': "'auto'",
+    \ 'use_environment': "'auto'",
+    \ 'environment_paths': '[]',
     \ 'smart_auto_mappings': 1,
     \ 'use_tag_stack': 1
 \ }
@@ -275,6 +277,14 @@ call jedi#init_python()  " Might throw an error.
 " ------------------------------------------------------------------------
 " functions that call python code
 " ------------------------------------------------------------------------
+function! jedi#use_environment(...) abort
+  PythonJedi jedi_vim.set_environment(*vim.eval('a:000'))
+endfunction
+
+function! jedi#use_environment_for_buffer(env) abort
+  PythonJedi jedi_vim.set_environment(*vim.eval('[env, bufnr("%")]'))
+endfunction
+
 function! jedi#goto() abort
     PythonJedi jedi_vim.goto(mode="goto")
 endfunction
@@ -624,6 +634,13 @@ function! jedi#setup_completion() abort
     endif
 endfunction
 
+" TODO: fallback to completing files (i.e. return known envs first, then
+" files?!).
+function! jedi#complete_environments(argl, cmdl, pos) abort
+    PythonJedi jedi_vim.complete_environments()
+EOF
+endfunction
+
 "PythonJedi jedi_vim.jedi.set_debug_function(jedi_vim.print_to_stdout, speed=True, warnings=False, notices=False)
 "PythonJedi jedi_vim.jedi.set_debug_function(jedi_vim.print_to_stdout)
 

plugin/jedi.vim

@@ -63,6 +63,22 @@ function! s:jedi_debug_info()
     call jedi#debug_info()
 endfunction
 command! -nargs=0 -bar JediDebugInfo call s:jedi_debug_info()
+
+" Transfer bang into current buffer number.
+" TODO: use window instead!?
+function! s:use_environment(bang, env) abort
+  let bufnr = a:bang ? +bufnr('%') : 0
+  if a:env ==# 'auto'
+    let env = ''
+  elseif a:env[-4:] ==# ' (*)'
+    let env = a:env[:-5]
+  else
+    let env = a:env
+  endif
+  return jedi#use_environment(env, bufnr)
+endfunction
+command! -bang -nargs=? -complete=custom,jedi#complete_environments JediUseEnvironment :call s:use_environment(<bang>0, <q-args>)
+
 command! -nargs=0 -bang JediClearCache call jedi#clear_cache(<bang>0)
 
 " vim: set et ts=4:

pythonx/jedi_vim.py

@@ -157,49 +157,84 @@ def wrapper(*args, **kwargs):
     return func_receiver
 
 
+# XXX: might use lru cache, given that you can configure it per buffer?!
 current_environment = (None, None)
 
 
+def set_environment(executable=None, bufnr=None):
+    global current_environment
+
+    # print(repr(executable), repr(bufnr))
+
+    if not executable:  # "unset"
+        executable = 'auto'
+        environment = None
+    else:
+        try:
+            # XXX: create_environment is not deterministic
+            #      (https://github.com/davidhalter/jedi/issues/1155)
+            # jedi.create_environment(executable, safe=False)
+            environment = jedi.api.environment.Environment(executable)
+        except jedi.InvalidPythonEnvironment as exc:
+            echo_highlight('executable=%s is not supported: %s.' % (
+                executable, str(exc)))
+            return
+    bufnr = int(bufnr)
+    if bufnr:
+        vim_command("call setbufvar(%d, 'jedi_use_environment', %r)" % (
+                    bufnr, executable))
+    else:
+        vim_command("let g:jedi#use_environment = %r" % executable)
+    current_environment = (executable, environment)
+
+
 def get_environment(use_cache=True):
     global current_environment
 
-    vim_force_python_version = vim_eval("g:jedi#force_py_version")
-    if use_cache and vim_force_python_version == current_environment[0]:
+    use_environment = vim_eval(
+        "get(b:, 'jedi_use_environment', get(g:, 'jedi#use_environment', ''))")
+    if use_cache and use_environment == current_environment[0]:
         return current_environment[1]
 
     environment = None
-    if vim_force_python_version == "auto":
+    if use_environment == "auto":
         environment = jedi.api.environment.get_cached_default_environment()
     else:
-        force_python_version = vim_force_python_version
-        if '0000' in force_python_version or '9999' in force_python_version:
-            # It's probably a float that wasn't shortened.
-            try:
-                force_python_version = "{:.1f}".format(float(force_python_version))
-            except ValueError:
-                pass
-        elif isinstance(force_python_version, float):
-            force_python_version = "{:.1f}".format(force_python_version)
-
         try:
-            environment = jedi.get_system_environment(force_python_version)
+            # XXX: create_environment is not deterministic
+            #      (https://github.com/davidhalter/jedi/issues/1155)
+            # jedi.create_environment(executable, safe=False)
+            environment = jedi.api.environment.Environment(use_environment)
         except jedi.InvalidPythonEnvironment as exc:
             environment = jedi.api.environment.get_cached_default_environment()
             echo_highlight(
-                "force_python_version=%s is not supported: %s - using %s." % (
-                    vim_force_python_version, str(exc), str(environment)))
+                'use_environment=%s is not supported: %s - using %s.' % (
+                    use_environment, str(exc), str(environment)))
 
-    current_environment = (vim_force_python_version, environment)
+    current_environment = (use_environment, environment)
     return environment
 
 
 def get_known_environments():
     """Get known Jedi environments."""
-    envs = list(jedi.api.environment.find_virtualenvs())
+    paths = vim_eval("g:jedi#environment_paths")
+    envs = list(jedi.api.environment.find_virtualenvs(paths=paths, safe=False))
     envs.extend(jedi.api.environment.find_system_environments())
     return envs
 
 
+def complete_environments():
+    envs = get_known_environments()
+    try:
+        current_executable = current_environment[1].executable
+    except AttributeError:
+        current_executable = None
+    vim.command("return '%s'" % '\n'.join(
+        ('%s (*)' if env.executable == current_executable else '%s') % (
+            env.executable
+        ) for env in envs))
+
+
 @catch_and_print_exceptions
 def get_script(source=None, column=None):
     jedi.settings.additional_dynamic_modules = [