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/cypress/support/webpack.cypress.ts

@@ -10,10 +10,17 @@ const buildConfig = () => {
       workerPublicPath: '/__cypress/src/',
     },
     devServer: {
-      static: {
-        directory: path.join(__dirname, '../../public'),
-        watch: false,
-      },
+      static: [
+        {
+          directory: path.join(__dirname, '../../public'),
+          watch: false,
+        },
+        {
+          directory: path.join(__dirname, '../fixtures/pyodide-packages'),
+          publicPath: '/pyodide-packages/',
+          watch: false,
+        },
+      ],
       port: 3200,
     },
     stats: 'none',

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(() => {

services/web/test/frontend/features/ide-react/components/python-output-pane.spec.tsx

@@ -298,4 +298,47 @@ describe('<PythonOutputPane />', function () {
       'ide-redesign-python-output-pane-line-info'
     )
   })
+
+  it('can load common python data analysis packages on code execution', function () {
+    const executablePythonFileContents = [
+      'import tomli',
+      '',
+      "print(tomli.loads('greeting = \"hello from tomli\"')['greeting'])",
+    ].join('\n')
+
+    const projectFiles = {
+      [pythonExecutableScript.filename]: executablePythonFileContents,
+    }
+    const ProjectProvider = makeProjectProvider(projectFiles)
+
+    cy.mount(
+      <EditorProviders
+        scope={{
+          editor: {
+            sharejs_doc: {
+              doc_id: pythonExecutableScript.file_id,
+              getSnapshot: () => executablePythonFileContents,
+            },
+            currentDocumentId: pythonExecutableScript.file_id,
+            openDocName: pythonExecutableScript.filename,
+          },
+        }}
+        providers={{ FileTreePathProvider, ProjectProvider }}
+      >
+        <PythonExecutionProvider
+          packageBaseUrl={`${window.location.origin}/pyodide-packages/`}
+        >
+          <PythonOutputPane />
+        </PythonExecutionProvider>
+      </EditorProviders>
+    )
+
+    cy.findByRole('button', { name: 'Run Python code' })
+      .should('not.be.disabled')
+      .click()
+    cy.findByText("ModuleNotFoundError: No module named 'tomli'").should(
+      'not.exist'
+    )
+    cy.findByText('hello from tomli').should('exist')
+  })
 })

services/web/test/frontend/features/ide-react/unit/editor/pyodide-worker-client.spec.ts

@@ -346,7 +346,11 @@ describe('PyodideWorkerClient', function () {
 
       newWorker.emitMessage({ type: 'listening' })
       expect(newWorker.postedMessages).to.deep.equal([
-        { type: 'init', baseAssetPath: BASE_ASSET_PATH },
+        {
+          type: 'init',
+          baseAssetPath: BASE_ASSET_PATH,
+          packageBaseUrl: undefined,
+        },
       ])
     })