Lite auto-load imported modules with pyodide.loadPackagesFromImports

.changeset/heavy-dogs-start.md

@@ -0,0 +1,8 @@
+---
+"@gradio/lite": minor
+"@gradio/wasm": minor
+"gradio": minor
+"website": minor
+---
+
+feat:Lite auto-load imported modules with `pyodide.loadPackagesFromImports`

js/_website/src/lib/components/DemosLite.svelte

@@ -198,7 +198,7 @@
 	let controller: {
 		run_code: (code: string) => Promise<void>;
 		install: (requirements: string[]) => Promise<void>;
-	};
+	} & EventTarget;
 
 	function debounce<T extends any[]>(
 		func: (...args: T) => Promise<unknown>,
@@ -251,6 +251,14 @@
 			debounced_run_code = debounce(controller.run_code, debounce_timeout);
 			debounced_install = debounce(controller.install, debounce_timeout);
 
+			controller.addEventListener("modules-auto-loaded", (event) => {
+				console.debug("Modules auto-loaded", event);
+				const packages = (event as CustomEvent).detail as { name: string }[];
+				const packageNames = packages.map((pkg) => pkg.name);
+				selected_demo.requirements =
+					selected_demo.requirements.concat(packageNames);
+			});
+
 			mounted = true;
 		} catch (error) {
 			console.error("Error loading Gradio Lite:", error);

js/lite/src/LiteIndex.svelte

@@ -4,7 +4,7 @@
 	import "@gradio/theme/pollen.css";
 	import "@gradio/theme/typography.css";
 
-	import { onDestroy, SvelteComponent } from "svelte";
+	import { onDestroy, SvelteComponent, createEventDispatcher } from "svelte";
 	import Index from "@self/spa";
 	import Playground from "./Playground.svelte";
 	import ErrorDisplay from "./ErrorDisplay.svelte";
@@ -94,6 +94,12 @@
 		error = (event as CustomEvent).detail;
 	});
 
+	const dispatch = createEventDispatcher();
+
+	worker_proxy.addEventListener("modules-auto-loaded", (event) => {
+		dispatch("modules-auto-loaded", (event as CustomEvent).detail);
+	});
+
 	// Internally, the execution of `runPythonCode()` or `runPythonFile()` is queued
 	// and its promise will be resolved after the Pyodide is loaded and the worker initialization is done
 	// (see the await in the `onmessage` callback in the webworker code)

js/lite/src/dev/App.svelte

@@ -79,6 +79,12 @@ def hi(name):
 			playground: false,
 			layout: null
 		});
+		controller.addEventListener("modules-auto-loaded", (event) => {
+			const packages = (event as CustomEvent).detail as { name: string }[];
+			const packageNames = packages.map((pkg) => pkg.name);
+			requirements_txt +=
+				"\n" + packageNames.map((line) => line + "  # auto-loaded").join("\n");
+		});
 	});
 	onDestroy(() => {
 		controller.unmount();

js/lite/src/index.ts

@@ -21,18 +21,45 @@ import LiteIndex from "./LiteIndex.svelte";
 // As a result, the users of the Wasm app will have to load the CSS file manually.
 // const ENTRY_CSS = "__ENTRY_CSS__";
 
-export interface GradioAppController {
-	run_code: (code: string) => Promise<void>;
-	run_file: (path: string) => Promise<void>;
-	write: (
+export class GradioAppController extends EventTarget {
+	constructor(private lite_svelte_app: LiteIndex) {
+		super();
+
+		this.lite_svelte_app.$on("error", (event: CustomEvent) => {
+			this.dispatchEvent(new CustomEvent("error", { detail: event.detail }));
+		});
+		this.lite_svelte_app.$on("modules-auto-loaded", (event: CustomEvent) => {
+			this.dispatchEvent(
+				new CustomEvent("modules-auto-loaded", { detail: event.detail })
+			);
+		});
+	}
+
+	run_code = (code: string): Promise<void> => {
+		return this.lite_svelte_app.run_code(code);
+	};
+	run_file = (path: string): Promise<void> => {
+		return this.lite_svelte_app.run_file(path);
+	};
+	write = (
 		path: string,
 		data: string | ArrayBufferView,
 		opts: any
-	) => Promise<void>;
-	rename: (old_path: string, new_path: string) => Promise<void>;
-	unlink: (path: string) => Promise<void>;
-	install: (requirements: string[]) => Promise<void>;
-	unmount: () => void;
+	): Promise<void> => {
+		return this.lite_svelte_app.write(path, data, opts);
+	};
+	rename = (old_path: string, new_path: string): Promise<void> => {
+		return this.lite_svelte_app.rename(old_path, new_path);
+	};
+	unlink = (path: string): Promise<void> => {
+		return this.lite_svelte_app.unlink(path);
+	};
+	install = (requirements: string[]): Promise<void> => {
+		return this.lite_svelte_app.install(requirements);
+	};
+	unmount = (): void => {
+		this.lite_svelte_app.$destroy();
+	};
 }
 
 export interface Options {
@@ -88,17 +115,7 @@ export function create(options: Options): GradioAppController {
 		}
 	});
 
-	return {
-		run_code: app.run_code,
-		run_file: app.run_file,
-		write: app.write,
-		rename: app.rename,
-		unlink: app.unlink,
-		install: app.install,
-		unmount() {
-			app.$destroy();
-		}
-	};
+	return new GradioAppController(app);
 }
 
 /**

js/wasm/src/message-types.ts

@@ -1,4 +1,5 @@
 import type { ASGIScope } from "./asgi-types";
+import type { PackageData } from "pyodide";
 
 export interface EmscriptenFile {
 	data: string | ArrayBufferView;
@@ -113,4 +114,11 @@ export interface OutMessageProgressUpdate extends OutMessageBase {
 		log: string;
 	};
 }
-export type OutMessage = OutMessageProgressUpdate;
+export interface OutMessageModulesAutoLoaded extends OutMessageBase {
+	type: "modules-auto-loaded";
+	data: {
+		packages: PackageData[];
+	};
+}
+
+export type OutMessage = OutMessageProgressUpdate | OutMessageModulesAutoLoaded;

js/wasm/src/webworker/index.ts

@@ -2,6 +2,7 @@
 /* eslint-env worker */
 
 import type {
+	PackageData,
 	PyodideInterface,
 	loadPyodide as loadPyodideValue
 } from "pyodide";
@@ -177,7 +178,8 @@ anyio.to_thread.run_sync = mocked_anyio_to_thread_run_sync
 async function initializeApp(
 	appId: string,
 	options: InMessageInitApp["data"],
-	updateProgress: (log: string) => void
+	updateProgress: (log: string) => void,
+	onModulesAutoLoaded: (packages: PackageData[]) => void
 ): Promise<void> {
 	const appHomeDir = getAppHomeDir(appId);
 	console.debug("Creating a home directory for the app.", {
@@ -188,6 +190,7 @@ async function initializeApp(
 
 	console.debug("Mounting files.", options.files);
 	updateProgress("Mounting files");
+	const pythonFileContents: string[] = [];
 	await Promise.all(
 		Object.keys(options.files).map(async (path) => {
 			const file = options.files[path];
@@ -206,6 +209,10 @@ async function initializeApp(
 			const appifiedPath = resolveAppHomeBasedPath(appId, path);
 			console.debug(`Write a file "${appifiedPath}"`);
 			writeFileWithParents(pyodide, appifiedPath, data, opts);
+
+			if (typeof data === "string" && path.endsWith(".py")) {
+				pythonFileContents.push(data);
+			}
 		})
 	);
 	console.debug("Files are mounted.");
@@ -215,7 +222,22 @@ async function initializeApp(
 	await micropip.install.callKwargs(options.requirements, { keep_going: true });
 	console.debug("Packages are installed.");
 
-	if (options.requirements.includes("matplotlib")) {
+	console.debug("Auto-loading modules.");
+	const loadedPackagesArr = await Promise.all(
+		pythonFileContents.map((source) => pyodide.loadPackagesFromImports(source))
+	);
+	const loadedPackagesSet = new Set(loadedPackagesArr.flat()); // Remove duplicates
+	const loadedPackages = Array.from(loadedPackagesSet);
+	if (loadedPackages.length > 0) {
+		onModulesAutoLoaded(loadedPackages);
+	}
+	const loadedPackageNames = loadedPackages.map((pkg) => pkg.name);
+	console.debug("Modules are auto-loaded.", loadedPackages);
+
+	if (
+		options.requirements.includes("matplotlib") ||
+		loadedPackageNames.includes("matplotlib")
+	) {
 		console.debug("Setting matplotlib backend.");
 		updateProgress("Setting matplotlib backend");
 		// Ref: https://github.com/pyodide/pyodide/issues/561#issuecomment-1992613717
@@ -276,6 +298,15 @@ function setupMessageHandler(receiver: MessageTransceiver): void {
 		};
 		receiver.postMessage(message);
 	};
+	const onModulesAutoLoaded = (packages: PackageData[]) => {
+		const message: OutMessage = {
+			type: "modules-auto-loaded",
+			data: {
+				packages
+			}
+		};
+		receiver.postMessage(message);
+	};
 
 	// App initialization is per app or receiver, so its promise is managed in this scope.
 	let appReadyPromise: Promise<void> | undefined = undefined;
@@ -322,7 +353,12 @@ function setupMessageHandler(receiver: MessageTransceiver): void {
 			await envReadyPromise;
 
 			if (msg.type === "init-app") {
-				appReadyPromise = initializeApp(appId, msg.data, updateProgress);
+				appReadyPromise = initializeApp(
+					appId,
+					msg.data,
+					updateProgress,
+					onModulesAutoLoaded
+				);
 
 				const replyMessage: ReplyMessageSuccess = {
 					type: "reply:success",
@@ -349,6 +385,15 @@ function setupMessageHandler(receiver: MessageTransceiver): void {
 				case "run-python-code": {
 					unload_local_modules();
 
+					console.debug(`Auto install the requirements`);
+					const loadedPackages = await pyodide.loadPackagesFromImports(
+						msg.data.code
+					);
+					if (loadedPackages.length > 0) {
+						onModulesAutoLoaded(loadedPackages);
+					}
+					console.debug("Modules are auto-loaded.", loadedPackages);
+
 					await run_code(appId, getAppHomeDir(appId), msg.data.code);
 
 					const replyMessage: ReplyMessageSuccess = {
@@ -382,6 +427,16 @@ function setupMessageHandler(receiver: MessageTransceiver): void {
 				case "file:write": {
 					const { path, data: fileData, opts } = msg.data;
 
+					if (typeof fileData === "string" && path.endsWith(".py")) {
+						console.debug(`Auto install the requirements in ${path}`);
+						const loadedPackages =
+							await pyodide.loadPackagesFromImports(fileData);
+						if (loadedPackages.length > 0) {
+							onModulesAutoLoaded(loadedPackages);
+						}
+						console.debug("Modules are auto-loaded.", loadedPackages);
+					}
+
 					const appifiedPath = resolveAppHomeBasedPath(appId, path);
 
 					console.debug(`Write a file "${appifiedPath}"`);

js/wasm/src/worker-proxy.ts

@@ -166,6 +166,13 @@ export class WorkerProxy extends EventTarget {
 				);
 				break;
 			}
+			case "modules-auto-loaded": {
+				this.dispatchEvent(
+					new CustomEvent("modules-auto-loaded", {
+						detail: msg.data.packages
+					})
+				);
+			}
 		}
 	}