Jedi-based Python code completion on gr.Code (#10812)

* Jedi-based code completion

* Delete unused objects

* add changeset

* add changeset

* Format

---------

Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>

Author: whitphx

.changeset/silver-bears-draw.md

@@ -0,0 +1,7 @@
+---
+"@gradio/code": minor
+"@gradio/wasm": minor
+"gradio": minor
+---
+
+feat:Jedi-based Python code completion on `gr.Code`

js/code/shared/Code.svelte

@@ -10,11 +10,13 @@
 	import { StateEffect, EditorState, type Extension } from "@codemirror/state";
 	import { indentWithTab } from "@codemirror/commands";
 	import { autocompletion, acceptCompletion } from "@codemirror/autocomplete";
+	import { LanguageSupport } from "@codemirror/language";
 
 	import { basicDark } from "cm6-theme-basic-dark";
 	import { basicLight } from "cm6-theme-basic-light";
 	import { basicSetup } from "./extensions";
 	import { getLanguageExtension } from "./language";
+	import { create_pyodide_autocomplete } from "./autocomplete";
 
 	export let class_names = "";
 	export let value = "";
@@ -42,8 +44,19 @@
 
 	$: get_lang(language);
 
+	const pyodide_autocomplete = create_pyodide_autocomplete();
+
 	async function get_lang(val: string): Promise<void> {
 		const ext = await getLanguageExtension(val);
+		if (
+			pyodide_autocomplete &&
+			val === "python" &&
+			ext instanceof LanguageSupport
+		) {
+			(ext.support as Extension[]).push(
+				ext.language.data.of({ autocomplete: pyodide_autocomplete })
+			);
+		}
 		lang_extension = ext;
 	}
 

js/code/shared/autocomplete.ts

@@ -0,0 +1,63 @@
+import type {
+	CompletionContext,
+	CompletionResult
+} from "@codemirror/autocomplete";
+import { getWorkerProxyContext } from "@gradio/wasm/svelte";
+import type { WorkerProxy } from "@gradio/wasm";
+
+// Jedi's completion types to CodeMirror's completion types.
+// If not defined here, Jedi's completion types will be used.
+const completion_type_map: Record<string, string> = {
+	module: "namespace"
+};
+
+type CodeMirrorAutocompleteAsyncFn = (
+	context: CompletionContext
+) => Promise<CompletionResult | null>;
+
+export function create_pyodide_autocomplete(): CodeMirrorAutocompleteAsyncFn | null {
+	let maybe_worker_proxy: WorkerProxy | undefined;
+	try {
+		maybe_worker_proxy = getWorkerProxyContext();
+	} catch (e) {
+		console.debug("Not in the Wasm env. Context-aware autocomplete disabled.");
+		return null;
+	}
+	if (!maybe_worker_proxy) {
+		return null;
+	}
+	const worker_proxy = maybe_worker_proxy;
+
+	return async function pyodide_autocomplete(
+		context: CompletionContext
+	): Promise<CompletionResult | null> {
+		try {
+			const completions = await worker_proxy.getCodeCompletions({
+				code: context.state.doc.toString(),
+				line: context.state.doc.lineAt(context.state.selection.main.head)
+					.number,
+				column:
+					context.state.selection.main.head -
+					context.state.doc.lineAt(context.state.selection.main.head).from
+			});
+			if (completions.length === 0) {
+				return null;
+			}
+			return {
+				from:
+					context.state.selection.main.head -
+					completions[0].completion_prefix_length,
+				options: completions.map((completion) => ({
+					label: completion.label,
+					type: completion_type_map[completion.type] ?? completion.type,
+					documentation: completion.docstring,
+					// Items starting with "_" are private attributes and should be sorted last.
+					boost: completion.label.startsWith("_") ? -1 : 0
+				}))
+			};
+		} catch (e) {
+			console.error("Error getting completions", e);
+			return null;
+		}
+	};
+}

js/code/shared/language.ts

@@ -1,26 +1,7 @@
 import type { Extension } from "@codemirror/state";
 import { StreamLanguage } from "@codemirror/language";
-import { sql } from "@codemirror/legacy-modes/mode/sql";
 import { _ } from "svelte-i18n";
 
-const possible_langs = [
-	"python",
-	"c",
-	"cpp",
-	"markdown",
-	"json",
-	"html",
-	"css",
-	"javascript",
-	"jinja2",
-	"typescript",
-	"yaml",
-	"dockerfile",
-	"shell",
-	"r",
-	"sql"
-];
-
 const sql_dialects = [
 	"standardSQL",
 	"msSQL",

js/wasm/src/message-types.ts

@@ -1,4 +1,5 @@
 import type { ASGIScope } from "./asgi-types";
+import type { CodeCompletionRequest } from "./webworker/code-completion/index";
 import type { PackageData } from "pyodide";
 
 export interface EmscriptenFile {
@@ -74,6 +75,10 @@ export interface InMessageInstall extends InMessageBase {
 		requirements: string[];
 	};
 }
+export interface InMessageCodeCompletion extends InMessageBase {
+	type: "code-completion";
+	data: CodeCompletionRequest;
+}
 
 export interface InMessageEcho extends InMessageBase {
 	// For debug
@@ -91,6 +96,7 @@ export type InMessage =
 	| InMessageFileRename
 	| InMessageFileUnlink
 	| InMessageInstall
+	| InMessageCodeCompletion
 	| InMessageEcho;
 
 export interface ReplyMessageSuccess<T = unknown> {

js/wasm/src/webworker/code-completion/code_completion.py

@@ -0,0 +1,16 @@
+import jedi
+
+
+def get_code_completions(code: str, line: int, column: int):
+    script = jedi.Script(code)
+    completions = script.complete(line, column)
+    serializable_completions = [
+        {
+            "label": completion.name,
+            "type": completion.type,
+            "docstring": completion.docstring(raw=True),
+            "completion_prefix_length": completion.get_completion_prefix_length(),
+        }
+        for completion in completions
+    ]
+    return serializable_completions

js/wasm/src/webworker/code-completion/index.ts

@@ -0,0 +1,70 @@
+import type { PyodideInterface } from "pyodide";
+import code_completion_py from "./code_completion.py?raw";
+import type { PyProxy } from "pyodide/ffi";
+
+export interface CodeCompletionRequest {
+	code: string;
+	// line is 1-based and column is 0-based due to Jedi's spec: https://jedi.readthedocs.io/en/latest/docs/api.html#jedi.Script
+	line: number; // 1-based
+	column: number; // 0-based
+}
+export interface CodeCompletion {
+	label: string;
+	type: string;
+	docstring: string;
+	completion_prefix_length: number;
+}
+export type CodeCompletionResponse = CodeCompletion[];
+
+type GetCodeCompletionsPyFn = (
+	code: string,
+	line: number,
+	column: number
+) => PyProxy;
+
+export class CodeCompleter {
+	private setupPromise: Promise<GetCodeCompletionsPyFn | null>;
+
+	constructor(private pyodide: PyodideInterface) {
+		// NOTE: obviously this constructor has a side effect on the Pyodide instance.
+		this.setupPromise = this.setup().catch((err) => {
+			console.error("Error while setting up code completion", err);
+			return null;
+		});
+	}
+
+	private async setup(): Promise<GetCodeCompletionsPyFn> {
+		const micropip = this.pyodide.pyimport("micropip");
+		await micropip.install.callKwargs(["jedi"], {
+			keep_going: true
+		});
+
+		this.pyodide.runPython(code_completion_py);
+
+		return this.pyodide.globals.get(
+			"get_code_completions"
+		) as GetCodeCompletionsPyFn;
+	}
+
+	public async getCodeCompletions(
+		request: CodeCompletionRequest
+	): Promise<CodeCompletionResponse> {
+		const getCodeCompletionsPyFn = await this.setupPromise;
+		if (!getCodeCompletionsPyFn) {
+			// Setting up failed, return empty response
+			console.debug("Code completion setup failed, returning empty response");
+			return [];
+		}
+
+		const { code, line, column } = request;
+		const completionsPy = getCodeCompletionsPyFn(code, line, column);
+		const completions: CodeCompletionResponse = completionsPy.toJs({
+			dict_converter: Object.fromEntries // dict -> object
+		});
+		// > ... If the return value is a PyProxy, you must explicitly destroy it or else it will be leaked.
+		// https://pyodide.org/en/stable/usage/type-conversions.html#calling-python-objects-from-javascript
+		completionsPy.destroy();
+
+		return completions;
+	}
+}

js/wasm/src/webworker/index.ts

@@ -24,16 +24,14 @@ import {
 import { patchRequirements, verifyRequirements } from "./requirements";
 import { makeAsgiRequest } from "./asgi";
 import { generateRandomString } from "./random";
+import { CodeCompleter } from "./code-completion";
 import scriptRunnerPySource from "./py/script_runner.py?raw";
 import unloadModulesPySource from "./py/unload_modules.py?raw";
 
 importScripts("https://cdn.jsdelivr.net/pyodide/v0.27.3/full/pyodide.js");
 
 type MessageTransceiver = DedicatedWorkerGlobalScope | MessagePort;
 
-let pyodide: PyodideInterface;
-let micropip: PyProxy;
-
 declare let loadPyodide: typeof loadPyodideValue; // This will be dynamically loaded by importScript.
 
 let call_asgi_app_from_js: (
@@ -56,6 +54,8 @@ let run_script: (
 let unload_local_modules: (target_dir_path?: string) => void;
 
 async function installPackages(
+	pyodide: PyodideInterface,
+	micropip: PyProxy,
 	requirements: string[],
 	retries = 3
 ): Promise<void> {
@@ -80,15 +80,20 @@ async function installPackages(
 	}
 }
 
+interface GradioLitePyodideEnvironment {
+	pyodide: PyodideInterface;
+	micropip: PyProxy;
+	codeCompleter: CodeCompleter;
+}
 async function initializeEnvironment(
 	options: InMessageInitEnv["data"],
 	updateProgress: (log: string) => void,
 	stdout: (output: string) => void,
 	stderr: (output: string) => void
-): Promise<void> {
+): Promise<GradioLitePyodideEnvironment> {
 	console.debug("Loading Pyodide.");
 	updateProgress("Loading Pyodide");
-	pyodide = await loadPyodide({
+	const pyodide = await loadPyodide({
 		stdout,
 		stderr
 	});
@@ -97,7 +102,7 @@ async function initializeEnvironment(
 	console.debug("Loading micropip");
 	updateProgress("Loading micropip");
 	await pyodide.loadPackage("micropip");
-	micropip = pyodide.pyimport("micropip");
+	const micropip = pyodide.pyimport("micropip");
 	console.debug("micropip is loaded.");
 
 	const gradioWheelUrls = [
@@ -108,7 +113,7 @@ async function initializeEnvironment(
 	updateProgress("Loading Gradio wheels");
 	await pyodide.loadPackage(["ssl", "setuptools"]);
 	await micropip.add_mock_package("ffmpy", "0.3.0");
-	await installPackages(gradioWheelUrls);
+	await installPackages(pyodide, micropip, gradioWheelUrls);
 	console.debug("Gradio wheels are loaded.");
 
 	console.debug("Mocking os module methods.");
@@ -196,9 +201,19 @@ anyio.to_thread.run_sync = mocked_anyio_to_thread_run_sync
 	console.debug("Python utility functions are set up.");
 
 	updateProgress("Initialization completed");
+
+	const codeCompleter = new CodeCompleter(pyodide);
+
+	return {
+		pyodide,
+		micropip,
+		codeCompleter
+	};
 }
 
 async function initializeApp(
+	pyodide: PyodideInterface,
+	micropip: PyProxy,
 	appId: string,
 	options: InMessageInitApp["data"],
 	updateProgress: (log: string) => void,
@@ -242,7 +257,7 @@ async function initializeApp(
 
 	console.debug("Installing packages.", options.requirements);
 	updateProgress("Installing packages");
-	await installPackages(options.requirements);
+	await installPackages(pyodide, micropip, options.requirements);
 	console.debug("Packages are installed.");
 
 	console.debug("Auto-loading modules.");
@@ -298,7 +313,8 @@ if ("postMessage" in ctx) {
 }
 
 // Environment initialization is global and should be done only once, so its promise is managed in a global scope.
-let envReadyPromise: Promise<void> | undefined = undefined;
+let envReadyPromise: Promise<GradioLitePyodideEnvironment> | undefined =
+	undefined;
 
 function setupMessageHandler(receiver: MessageTransceiver): void {
 	// A concept of "app" is introduced to support multiple apps in a single worker.
@@ -404,13 +420,15 @@ function setupMessageHandler(receiver: MessageTransceiver): void {
 			if (envReadyPromise == null) {
 				throw new Error("Pyodide Initialization is not started.");
 			}
-			await envReadyPromise;
+			const { pyodide, micropip, codeCompleter } = await envReadyPromise;
 
 			const gradio = pyodide.pyimport("gradio");
 			gradio.wasm_utils.register_error_traceback_callback(appId, onPythonError);
 
 			if (msg.type === "init-app") {
 				appReadyPromise = initializeApp(
+					pyodide,
+					micropip,
 					appId,
 					msg.data,
 					updateProgress,
@@ -541,7 +559,7 @@ function setupMessageHandler(receiver: MessageTransceiver): void {
 
 					console.debug("Install the requirements:", requirements);
 					verifyRequirements(requirements); // Blocks the not allowed wheel URL schemes.
-					await installPackages(requirements)
+					await installPackages(pyodide, micropip, requirements)
 						.then(() => {
 							if (requirements.includes("matplotlib")) {
 								// Ref: https://github.com/pyodide/pyodide/issues/561#issuecomment-1992613717
@@ -566,6 +584,16 @@ except ImportError:
 						});
 					break;
 				}
+				case "code-completion": {
+					const request = msg.data;
+					const completions = await codeCompleter.getCodeCompletions(request);
+					const replyMessage: ReplyMessageSuccess = {
+						type: "reply:success",
+						data: completions
+					};
+					messagePort.postMessage(replyMessage);
+					break;
+				}
 			}
 		} catch (error) {
 			console.error(error);