Allow to use Jedi language server in system environment or virtual environments

As this extension uses an internal Jedi LSP, it is hard for the users to
change the jedi setting to adapt various requirements. And in some cases,
the user require a specific version of the Jedi, which is hard to change
as it is embedded.

Add a new configuration property to allow the users to use Jedi LSP in
their system environment or virtual environments, and auto fallback to
the internal version if the external version is not found.

Signed-off-by: Inochi Amaoto <inochiama@gmail.com>

Author: inochisa

package.json

@@ -536,6 +536,12 @@
                     "scope": "window",
                     "type": "string"
                 },
+                "python.jedi.useJediInEnvPath": {
+                    "default": false,
+                    "description": "%python.jedi.useJediInEnvPath.description%",
+                    "scope": "window",
+                    "type": "boolean"
+                },
                 "python.interpreter.infoVisibility": {
                     "default": "onPythonRelated",
                     "description": "%python.interpreter.infoVisibility.description%",

package.nls.json

@@ -56,6 +56,7 @@
     "python.languageServer.jediDescription": "Use Jedi behind the Language Server Protocol (LSP) as a language server.",
     "python.languageServer.pylanceDescription": "Use Pylance as a language server.",
     "python.languageServer.noneDescription": "Disable language server capabilities.",
+    "python.jedi.useJediInEnvPath.description": "Use Jedi in system environment (if pyenv is disabled) or virtual environments folder (if pyenv is enabled) instead of the builtin Jedi.",
     "python.interpreter.infoVisibility.description": "Controls when to display information of selected interpreter in the status bar.",
     "python.interpreter.infoVisibility.never.description": "Never display information.",
     "python.interpreter.infoVisibility.onPythonRelated.description": "Only display information if Python-related files are opened.",

python_files/run-jedi-language-server.py

@@ -2,11 +2,21 @@
 import pathlib
 import sys
 
-# Add the lib path to our sys path so jedi_language_server can find its references
-extension_dir = pathlib.Path(__file__).parent.parent
-EXTENSION_ROOT = os.fsdecode(extension_dir)
-sys.path.insert(0, os.fsdecode(extension_dir / "python_files" / "lib" / "jedilsp"))
-del extension_dir
+use_external_jedi = os.getenv("USE_JEDI_IN_ENV", '0') == "1"
+
+if use_external_jedi:
+    try:
+        import jedi_language_server
+    except Exception:
+        print("External Jedi is not found, fallback to the builtin one.", file=sys.stderr)
+        use_external_jedi = False
+
+if not use_external_jedi:
+    # Add the lib path to our sys path so jedi_language_server can find its references
+    extension_dir = pathlib.Path(__file__).parent.parent
+    EXTENSION_ROOT = os.fsdecode(extension_dir)
+    sys.path.insert(0, os.fsdecode(extension_dir / "python_files" / "lib" / "jedilsp"))
+    del extension_dir
 
 
 from jedi_language_server.cli import cli  # noqa: E402

src/client/activation/jedi/languageClientFactory.ts

@@ -2,7 +2,8 @@
 // Licensed under the MIT License.
 
 import * as path from 'path';
-import { LanguageClient, LanguageClientOptions, ServerOptions } from 'vscode-languageclient/node';
+import { WorkspaceConfiguration } from 'vscode';
+import { ExecutableOptions, LanguageClient, LanguageClientOptions, ServerOptions } from 'vscode-languageclient/node';
 
 import { EXTENSION_ROOT_DIR, PYTHON_LANGUAGE } from '../../common/constants';
 import { Resource } from '../../common/types';
@@ -13,7 +14,10 @@ import { ILanguageClientFactory } from '../types';
 const languageClientName = 'Python Jedi';
 
 export class JediLanguageClientFactory implements ILanguageClientFactory {
-    constructor(private interpreterService: IInterpreterService) {}
+    constructor(
+        private interpreterService: IInterpreterService,
+        private readonly workspaceConfiguration: WorkspaceConfiguration,
+    ) {}
 
     public async createLanguageClient(
         resource: Resource,
@@ -23,9 +27,16 @@ export class JediLanguageClientFactory implements ILanguageClientFactory {
         // Just run the language server using a module
         const lsScriptPath = path.join(EXTENSION_ROOT_DIR, 'python_files', 'run-jedi-language-server.py');
         const interpreter = await this.interpreterService.getActiveInterpreter(resource);
+        const useJediInEnv = this.workspaceConfiguration.get<boolean>('jedi.useJediInEnvPath') === true;
+        const envVars: NodeJS.ProcessEnv = {
+            USE_JEDI_IN_ENV: useJediInEnv ? '1' : '0',
+            ...process.env,
+        };
+        const executableOptions: ExecutableOptions = { env: envVars };
         const serverOptions: ServerOptions = {
             command: interpreter ? interpreter.path : 'python',
             args: [lsScriptPath],
+            options: executableOptions,
         };
 
         return new LanguageClient(PYTHON_LANGUAGE, languageClientName, serverOptions, clientOptions);

src/client/languageServer/jediLSExtensionManager.ts

@@ -47,7 +47,10 @@ export class JediLSExtensionManager implements IDisposable, ILanguageServerExten
             configurationService,
             workspaceService,
         );
-        this.clientFactory = new JediLanguageClientFactory(interpreterService);
+        this.clientFactory = new JediLanguageClientFactory(
+            interpreterService,
+            workspaceService.getConfiguration('python'),
+        );
         this.serverProxy = new JediLanguageServerProxy(this.clientFactory);
         this.serverManager = new JediLanguageServerManager(
             serviceContainer,