Merge pull request #32943 from overleaf/cd-auto-install-python-packages

Auto-install python packages from the executing python script

GitOrigin-RevId: e343312d61e1804d927688bf4e0de00b2bdb5382

services/web/frontend/js/features/ide-react/components/editor/python/pyodide-worker-client.ts

@@ -30,6 +30,7 @@ export type LifecycleCallback = (
 export class PyodideWorkerClient {
   private worker: Worker
   private baseAssetPath: string
+  private packageBaseUrl: string | undefined
   private createWorker: () => Worker
   private listening = false
   private destroyed = false
@@ -40,11 +41,13 @@ export class PyodideWorkerClient {
 
   constructor(options: {
     baseAssetPath: string
+    packageBaseUrl?: string
     createWorker: () => Worker
     onOutput?: OutputCallback
     onLifecycle?: LifecycleCallback
   }) {
     this.baseAssetPath = options.baseAssetPath
+    this.packageBaseUrl = options.packageBaseUrl
     this.createWorker = options.createWorker
     this.outputCallback = options.onOutput ?? null
     this.lifecycleCallback = options.onLifecycle ?? null
@@ -54,6 +57,7 @@ export class PyodideWorkerClient {
     this.queueMessage({
       type: 'init',
       baseAssetPath: this.baseAssetPath,
+      packageBaseUrl: this.packageBaseUrl,
     })
   }
 
@@ -97,6 +101,7 @@ export class PyodideWorkerClient {
     this.queueMessage({
       type: 'init',
       baseAssetPath: this.baseAssetPath,
+      packageBaseUrl: this.packageBaseUrl,
     })
   }
 

services/web/frontend/js/features/ide-react/components/editor/python/pyodide-worker-messages.ts

@@ -15,6 +15,7 @@ export type OutputFileData = {
 export type InitRequest = {
   type: 'init'
   baseAssetPath: string
+  packageBaseUrl?: string
 }
 
 export type RunCodeRequest = {

services/web/frontend/js/features/ide-react/components/editor/python/pyodide.worker.ts

@@ -3,6 +3,7 @@ import path from 'path-browserify'
 import type { PyodideInterface } from 'pyodide'
 import type {
   OutputFileData,
+  InitRequest,
   ProjectFileData,
   PyodideWorkerRequest,
   RunCodeRequest,
@@ -14,6 +15,7 @@ type PyodideModule = typeof import('pyodide')
 const PROJECT_FS_ROOT = '/project'
 const PROJECT_FS_PREFIX = `${PROJECT_FS_ROOT}/`
 const PYODIDE_INDEX_PATH = 'js/libs/pyodide/'
+const PYODIDE_CDN_URL = 'https://cdn.jsdelivr.net/pyodide/v'
 
 function ensureDirectoryExists(fs: PyodideFS, filePath: string) {
   const directory = path.dirname(filePath)
@@ -49,6 +51,7 @@ function syncProjectFiles(fs: PyodideFS, files: ProjectFileData[]) {
 }
 
 let pyodideModule: PyodideModule | null = null
+let packageBaseUrlOverride: string | undefined
 
 async function loadPyodideModule(pyodideIndexUrl: string) {
   const runtimeModuleUrl = `${pyodideIndexUrl}pyodide.mjs`
@@ -66,12 +69,14 @@ async function loadPyodideModule(pyodideIndexUrl: string) {
   }
 }
 
-async function handleInit(msg: { baseAssetPath: string }) {
+async function handleInit(msg: InitRequest) {
   const pyodideIndexUrl = new URL(
     PYODIDE_INDEX_PATH,
     msg.baseAssetPath
   ).toString()
 
+  packageBaseUrlOverride = msg.packageBaseUrl
+
   try {
     pyodideModule = await loadPyodideModule(pyodideIndexUrl)
     self.postMessage({ type: 'loaded' })
@@ -109,6 +114,9 @@ async function handleRunCode(msg: RunCodeRequest) {
 
   const instance = await pyodideModule.loadPyodide({
     env: { MPLBACKEND: 'Agg' },
+    packageBaseUrl:
+      packageBaseUrlOverride ??
+      `${PYODIDE_CDN_URL}${pyodideModule.version}/full/`,
   })
 
   const writtenPaths = new Set<string>()
@@ -157,6 +165,7 @@ async function handleRunCode(msg: RunCodeRequest) {
       return originalWrite.call(fs, ...args)
     }) as PyodideFS['write']
 
+    await instance.loadPackagesFromImports(msg.code)
     const result = await instance.runPythonAsync(msg.code)
     if (result !== undefined) {
       self.postMessage({

services/web/frontend/js/features/ide-react/components/editor/python/python-runner.ts

@@ -43,6 +43,7 @@ export class PythonRunner {
   readonly fileId: string
   private client: PyodideWorkerClient | null = null
   private readonly baseAssetPath: string
+  private readonly packageBaseUrl: string | undefined
   private readonly createWorker: () => Worker
   private readonly getExecutionContext: () => Promise<ExecutionContext | null>
   private listeners = new Set<Listener>()
@@ -54,10 +55,12 @@ export class PythonRunner {
     fileId: string,
     baseAssetPath: string,
     getExecutionContext: () => Promise<ExecutionContext | null>,
-    createWorker: () => Worker
+    createWorker: () => Worker,
+    packageBaseUrl?: string
   ) {
     this.fileId = fileId
     this.baseAssetPath = baseAssetPath
+    this.packageBaseUrl = packageBaseUrl
     this.createWorker = createWorker
     this.getExecutionContext = getExecutionContext
   }
@@ -99,6 +102,7 @@ export class PythonRunner {
 
     this.client = new PyodideWorkerClient({
       baseAssetPath: this.baseAssetPath,
+      packageBaseUrl: this.packageBaseUrl,
       createWorker: this.createWorker,
       onLifecycle: event => {
         switch (event.type) {

services/web/frontend/js/features/ide-react/context/python-execution-context.tsx

@@ -37,9 +37,9 @@ export const PythonExecutionContext = createContext<
   PythonExecutionContextValue | undefined
 >(undefined)
 
-export const PythonExecutionProvider: FC<PropsWithChildren> = ({
-  children,
-}) => {
+export const PythonExecutionProvider: FC<
+  PropsWithChildren<{ packageBaseUrl?: string }>
+> = ({ children, packageBaseUrl }) => {
   const { openDocs } = useEditorManagerContext()
   const { projectSnapshot } = useProjectContext()
   const { pathInFolder } = useFileTreePathContext()
@@ -99,13 +99,14 @@ export const PythonExecutionProvider: FC<PropsWithChildren> = ({
         fileId,
         baseAssetPathRef.current,
         () => getExecutionContext(fileId),
-        createPyodideWorker
+        createPyodideWorker,
+        packageBaseUrl
       )
       runner.init()
       runnersRef.current.set(fileId, runner)
       return runner
     },
-    [getExecutionContext]
+    [getExecutionContext, packageBaseUrl]
   )
 
   useEffect(() => {