✨ feat: add builtin Python plugin (#8873)

* feat: 初步完成

* chore: type

* feat: 图片功能

* feat: 文件下载功能

* refactor: 简化代码

* chore: 清理代码

* chore: clean

* chore: 清理代码

* chore: 清理代码

* chore: 小改进

* fix: 上传完成前图片无法显示

* refactor: 增加 python-interpreter package

* chore: 清理

* feat: 传入上下文中的文件

* chore: 小优化

* chore: 中文字体

* chore: clean

* fix: 服务端部署

* fix: 重复文件检查

* test: 增加 interpreter.test.ts

* test: add worker.test.ts

* style: fix import

* test: fix

* style: fix import

* style: move env file to envs

* style: 限制代码框高度

* style: 重命名

* misc: 小修小补

* refactor: 重命名为 code-interpreter

---------

Co-authored-by: Arvin Xu <arvinx@foxmail.com>

packages/python-interpreter/package.json

@@ -0,0 +1,15 @@
+{
+  "name": "@lobechat/python-interpreter",
+  "version": "1.0.0",
+  "private": true,
+  "main": "src/index.ts",
+  "scripts": {
+    "test": "vitest",
+    "test:coverage": "vitest --coverage"
+  },
+  "dependencies": {
+    "comlink": "^4.4.2",
+    "pyodide": "^0.28.2",
+    "url-join": "^5.0.0"
+  }
+}

packages/python-interpreter/src/__tests__/interpreter.test.ts

@@ -0,0 +1,27 @@
+// @vitest-environment node
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+describe('Python interpreter', () => {
+  beforeEach(() => {
+    vi.resetModules();
+  });
+
+  it('should be undefined if is not in browser', async () => {
+    const { PythonInterpreter } = await import('../index');
+    expect(PythonInterpreter).toBeUndefined();
+  });
+
+  it('should be defined if is in browser', async () => {
+    const MockWorker = vi.fn().mockImplementation(() => ({
+      postMessage: vi.fn(),
+      terminate: vi.fn(),
+      addEventListener: vi.fn(),
+      removeEventListener: vi.fn(),
+    }));
+
+    vi.stubGlobal('Worker', MockWorker);
+
+    const { PythonInterpreter } = await import('../index');
+    expect(PythonInterpreter).toBeDefined();
+  });
+});

packages/python-interpreter/src/__tests__/worker.test.ts

@@ -0,0 +1,252 @@
+// @vitest-environment node
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+vi.mock('comlink', () => ({
+  expose: vi.fn(),
+}));
+
+describe('PythonWorker', () => {
+  const mockPyodide = {
+    FS: {
+      mkdirTree: vi.fn(),
+      chdir: vi.fn(),
+      writeFile: vi.fn(),
+      readdir: vi.fn(),
+      readFile: vi.fn(),
+    },
+    loadPackage: vi.fn(),
+    pyimport: vi.fn(),
+    loadPackagesFromImports: vi.fn(),
+    setStdout: vi.fn(),
+    setStderr: vi.fn(),
+    runPythonAsync: vi.fn(),
+    loadedPackages: {},
+  };
+
+  const mockMicropip = {
+    set_index_urls: vi.fn(),
+    install: vi.fn(),
+  };
+
+  beforeEach(() => {
+    vi.resetModules();
+    vi.clearAllMocks();
+
+    // Setup minimal global mocks
+    vi.stubGlobal('importScripts', vi.fn());
+    vi.stubGlobal('loadPyodide', vi.fn().mockResolvedValue(mockPyodide));
+    vi.stubGlobal(
+      'fetch',
+      vi.fn().mockResolvedValue({
+        arrayBuffer: () => Promise.resolve(new ArrayBuffer(1024)),
+      }),
+    );
+
+    mockPyodide.pyimport.mockReturnValue(mockMicropip);
+    mockPyodide.loadedPackages = {};
+  });
+
+  const importWorker = async () => {
+    const { PythonWorker } = await import('../worker');
+    return { PythonWorker };
+  };
+
+  describe('constructor', () => {
+    it('should initialize with default options', () => {
+      return importWorker().then(({ PythonWorker }) => {
+        const worker = new PythonWorker({});
+
+        expect(worker.pypiIndexUrl).toBe('PYPI');
+        expect(worker.pyodideIndexUrl).toBe('https://cdn.jsdelivr.net/pyodide/v0.28.2/full');
+        expect(worker.uploadedFiles).toEqual([]);
+      });
+    });
+
+    it('should initialize with custom options', () => {
+      const options = {
+        pyodideIndexUrl: 'https://test.cdn.com/pyodide',
+        pypiIndexUrl: 'https://test.pypi.org',
+      };
+      return importWorker().then(({ PythonWorker }) => {
+        const worker = new PythonWorker(options);
+
+        expect(worker.pypiIndexUrl).toBe('https://test.pypi.org');
+        expect(worker.pyodideIndexUrl).toBe('https://test.cdn.com/pyodide');
+      });
+    });
+
+    it('should call importScripts with pyodide.js', () => {
+      return importWorker().then(({ PythonWorker }) => {
+        new PythonWorker({});
+        expect(globalThis.importScripts).toHaveBeenCalledWith(
+          expect.stringContaining('/pyodide.js'),
+        );
+      });
+    });
+  });
+
+  describe('pyodide getter', () => {
+    it('should throw error when pyodide is not initialized', () => {
+      return importWorker().then(({ PythonWorker }) => {
+        const worker = new PythonWorker({});
+        expect(() => worker.pyodide).toThrow('Python interpreter not initialized');
+      });
+    });
+
+    it('should return pyodide when initialized', async () => {
+      const { PythonWorker } = await importWorker();
+      const worker = new PythonWorker({});
+      await worker.init();
+      expect(worker.pyodide).toBe(mockPyodide);
+    });
+  });
+
+  describe('init', () => {
+    it('should initialize pyodide and setup filesystem', async () => {
+      const { PythonWorker } = await importWorker();
+      const worker = new PythonWorker({
+        pyodideIndexUrl: 'https://test.cdn.com/pyodide',
+      });
+
+      await worker.init();
+
+      expect(globalThis.loadPyodide).toHaveBeenCalledWith({
+        indexURL: 'https://test.cdn.com/pyodide',
+      });
+      expect(mockPyodide.FS.mkdirTree).toHaveBeenCalledWith('/mnt/data');
+      expect(mockPyodide.FS.chdir).toHaveBeenCalledWith('/mnt/data');
+    });
+  });
+
+  describe('file operations', () => {
+    let worker: any;
+
+    beforeEach(async () => {
+      const { PythonWorker } = await importWorker();
+      worker = new PythonWorker({});
+      await worker.init();
+    });
+
+    it('should upload files correctly', async () => {
+      const mockFile = new File(['test content'], 'test.txt', { type: 'text/plain' });
+
+      await worker.uploadFiles([mockFile]);
+
+      expect(mockPyodide.FS.writeFile).toHaveBeenCalledWith(
+        '/mnt/data/test.txt',
+        expect.any(Uint8Array),
+      );
+      expect(worker.uploadedFiles).toContain(mockFile);
+    });
+
+    it('should upload files with absolute path as-is', async () => {
+      const absFile = new File([Uint8Array.from([1, 2])], '/abs.txt');
+      await worker.uploadFiles([absFile]);
+      expect(mockPyodide.FS.writeFile).toHaveBeenCalledWith('/abs.txt', expect.any(Uint8Array));
+    });
+
+    it('should download new files from filesystem', async () => {
+      const mockFileContent = new Uint8Array([1, 2, 3, 4]);
+
+      mockPyodide.FS.readdir.mockReturnValue(['.', '..', 'output.txt']);
+      (mockPyodide.FS as any).readFile.mockReturnValue(mockFileContent);
+
+      const files = await worker.downloadFiles();
+
+      expect(files).toHaveLength(1);
+      expect(files[0].name).toBe('/mnt/data/output.txt');
+    });
+
+    it('should skip identical files in download (dedup)', async () => {
+      const same = new File([Uint8Array.from([7, 8])], 'same.txt');
+      await worker.uploadFiles([same]);
+
+      mockPyodide.FS.readdir.mockReturnValue(['.', '..', 'same.txt']);
+      (mockPyodide.FS as any).readFile.mockReturnValue(Uint8Array.from([7, 8]));
+
+      const files = await worker.downloadFiles();
+      expect(files).toHaveLength(0);
+    });
+  });
+
+  describe('runPython', () => {
+    let worker: any;
+
+    beforeEach(async () => {
+      const { PythonWorker } = await importWorker();
+      worker = new PythonWorker({});
+      await worker.init();
+    });
+
+    it('should execute python code successfully', async () => {
+      const code = 'print("Hello, World!")';
+      const expectedResult = 'Hello, World!';
+
+      mockPyodide.runPythonAsync.mockResolvedValue(expectedResult);
+
+      const result = await worker.runPython(code);
+
+      expect(result.success).toBe(true);
+      expect(result.result).toBe(expectedResult);
+      expect(mockPyodide.runPythonAsync).toHaveBeenCalledWith(code);
+    });
+
+    it('should call loadPackagesFromImports with code', async () => {
+      const code = 'print("x")';
+      mockPyodide.runPythonAsync.mockResolvedValue('x');
+      await worker.runPython(code);
+      expect(mockPyodide.loadPackagesFromImports).toHaveBeenCalledWith(code);
+    });
+
+    it('should handle python execution errors', async () => {
+      const error = new Error('SyntaxError: invalid syntax');
+      mockPyodide.runPythonAsync.mockRejectedValue(error);
+
+      const result = await worker.runPython('invalid code');
+
+      expect(result.success).toBe(false);
+      expect(result.output).toContainEqual({
+        data: 'SyntaxError: invalid syntax',
+        type: 'stderr',
+      });
+    });
+
+    it('should install packages using micropip', async () => {
+      const packages = ['numpy', 'pandas'];
+
+      await worker.installPackages(packages);
+
+      expect(mockPyodide.loadPackage).toHaveBeenCalledWith('micropip');
+      expect(mockMicropip.set_index_urls).toHaveBeenCalledWith([worker.pypiIndexUrl, 'PYPI']);
+      expect(mockMicropip.install).toHaveBeenCalledWith(packages);
+    });
+
+    it('should patch matplotlib when loaded', async () => {
+      mockPyodide.loadedPackages = { matplotlib: true } as any;
+      mockPyodide.runPythonAsync.mockResolvedValueOnce(undefined).mockResolvedValueOnce('ok');
+      const res = await worker.runPython('print(1)');
+      expect(res.success).toBe(true);
+      expect(mockPyodide.runPythonAsync).toHaveBeenNthCalledWith(
+        1,
+        expect.stringContaining('patch_matplotlib()'),
+      );
+    });
+
+    it('should write fonts into truetype directory before run', async () => {
+      mockPyodide.runPythonAsync.mockResolvedValue('ok');
+      await worker.runPython('print(1)');
+      expect(mockPyodide.FS.mkdirTree).toHaveBeenCalledWith('/usr/share/fonts/truetype');
+      expect(mockPyodide.FS.writeFile).toHaveBeenCalledWith(
+        expect.stringContaining('/usr/share/fonts/truetype/STSong.ttf'),
+        expect.any(Uint8Array),
+      );
+    });
+
+    it('should stringify non-string result', async () => {
+      mockPyodide.runPythonAsync.mockResolvedValue({ toString: () => '42' });
+      const r = await worker.runPython('1+41');
+      expect(r.success).toBe(true);
+      expect(r.result).toBe('42');
+    });
+  });
+});

packages/python-interpreter/src/index.ts

@@ -0,0 +1,2 @@
+export { PythonInterpreter } from './interpreter';
+export * from './types';

packages/python-interpreter/src/interpreter.ts

@@ -0,0 +1,13 @@
+import * as Comlink from 'comlink';
+
+import type { PythonWorkerType } from './worker';
+
+export const PythonInterpreter = (() => {
+  if (typeof Worker !== 'undefined') {
+    let worker = new Worker(new URL('worker.ts', import.meta.url), {
+      type: 'module',
+    });
+    return Comlink.wrap<PythonWorkerType>(worker);
+  }
+  return undefined;
+})();

packages/python-interpreter/src/types.ts

@@ -0,0 +1,23 @@
+export interface PythonOptions {
+  /**
+   * Pyodide CDN URL
+   */
+  pyodideIndexUrl?: string;
+  /**
+   * PyPI 索引 URL，要求支持 [JSON API](https://warehouse.pypa.io/api-reference/json.html)
+   *
+   * 默认值：`https://pypi.org/pypi/{package_name}/json`
+   */
+  pypiIndexUrl?: string;
+}
+
+export interface PythonOutput {
+  data: string;
+  type: 'stdout' | 'stderr';
+}
+
+export interface PythonResult {
+  output?: PythonOutput[];
+  result?: string;
+  success: boolean;
+}

packages/python-interpreter/src/worker.ts

@@ -0,0 +1,211 @@
+import * as Comlink from 'comlink';
+import { PyodideAPI, loadPyodide as loadPyodideType } from 'pyodide';
+import urlJoin from 'url-join';
+
+import { PythonOptions, PythonOutput, PythonResult } from './types';
+
+declare global {
+  // eslint-disable-next-line no-var
+  var loadPyodide: typeof loadPyodideType;
+}
+
+const PATCH_MATPLOTLIB = `
+def patch_matplotlib():
+  import matplotlib
+  import matplotlib.pyplot as plt
+  from matplotlib import font_manager
+
+  # patch plt.show
+  matplotlib.use('Agg')
+  index = 1
+  def show():
+    nonlocal index
+    plt.savefig(f'/mnt/data/plot_{index}.png', format="png")
+    plt.clf()
+    index += 1
+  plt.show = show
+
+  # patch fonts
+  font_path = '/usr/share/fonts/truetype/STSong.ttf'
+  font_manager.fontManager.addfont(font_path)
+  plt.rcParams['font.family'] = 'STSong'
+
+patch_matplotlib()`;
+
+// Pyodide 对象不能在 Worker 之间传递，因此定义为全局变量
+let pyodide: PyodideAPI | undefined;
+
+class PythonWorker {
+  pyodideIndexUrl: string;
+  pypiIndexUrl: string;
+  uploadedFiles: File[];
+
+  constructor(options: PythonOptions) {
+    this.pypiIndexUrl = options.pypiIndexUrl || 'PYPI';
+    this.pyodideIndexUrl =
+      options.pyodideIndexUrl || 'https://cdn.jsdelivr.net/pyodide/v0.28.2/full';
+    globalThis.importScripts(urlJoin(this.pyodideIndexUrl, 'pyodide.js'));
+    this.uploadedFiles = [];
+  }
+
+  get pyodide() {
+    if (!pyodide) {
+      throw new Error('Python interpreter not initialized');
+    }
+    return pyodide;
+  }
+
+  /**
+   * 初始化 Python 解释器
+   */
+  async init() {
+    pyodide = await globalThis.loadPyodide({
+      indexURL: this.pyodideIndexUrl,
+    });
+    pyodide.FS.mkdirTree('/mnt/data');
+    pyodide.FS.chdir('/mnt/data');
+  }
+
+  /**
+   * 上传文件到解释器环境中
+   * @param files 文件列表
+   */
+  async uploadFiles(files: File[]) {
+    for (const file of files) {
+      const content = new Uint8Array(await file.arrayBuffer());
+      // TODO: 此处可以考虑使用 WORKERFS 减少一次拷贝
+      if (file.name.startsWith('/')) {
+        this.pyodide.FS.writeFile(file.name, content);
+      } else {
+        this.pyodide.FS.writeFile(`/mnt/data/${file.name}`, content);
+      }
+      this.uploadedFiles.push(file);
+    }
+  }
+
+  /**
+   * 从解释器环境中下载变动的文件
+   * @param files 文件列表
+   */
+  async downloadFiles() {
+    const result: File[] = [];
+    for (const entry of this.pyodide.FS.readdir('/mnt/data')) {
+      if (entry === '.' || entry === '..') continue;
+      const filePath = `/mnt/data/${entry}`;
+      // pyodide 的 FS 类型定义有问题，只能采用 any
+      const content = (this.pyodide.FS as any).readFile(filePath, { encoding: 'binary' });
+      const blob = new Blob([content]);
+      const file = new File([blob], filePath);
+      if (await this.isNewFile(file)) {
+        result.push(file);
+      }
+    }
+    return result;
+  }
+
+  /**
+   * 安装 Python 包
+   * @param packages 包名列表
+   */
+  async installPackages(packages: string[]) {
+    await this.pyodide.loadPackage('micropip');
+    const micropip = this.pyodide.pyimport('micropip');
+    micropip.set_index_urls([this.pypiIndexUrl, 'PYPI']);
+    await micropip.install(packages);
+  }
+
+  /**
+   * 执行 Python 代码
+   * @param code 代码
+   */
+  async runPython(code: string): Promise<PythonResult> {
+    await this.patchFonts();
+    // NOTE: loadPackagesFromImports 只会处理 pyodide 官方包
+    await this.pyodide.loadPackagesFromImports(code);
+    await this.patchPackages();
+
+    // 安装依赖后再捕获标准输出，避免记录安装日志
+    const output: PythonOutput[] = [];
+    this.pyodide.setStdout({
+      batched: (o: string) => {
+        output.push({ data: o, type: 'stdout' });
+      },
+    });
+    this.pyodide.setStderr({
+      batched: (o: string) => {
+        output.push({ data: o, type: 'stderr' });
+      },
+    });
+
+    // 执行代码
+    let result;
+    let success = false;
+    try {
+      result = await this.pyodide.runPythonAsync(code);
+      success = true;
+    } catch (error) {
+      output.push({
+        data: error instanceof Error ? error.message : String(error),
+        type: 'stderr',
+      });
+    }
+
+    return {
+      output,
+      result: result?.toString(),
+      success,
+    };
+  }
+
+  private async patchPackages() {
+    const hasMatplotlib = Object.keys(this.pyodide.loadedPackages).includes('matplotlib');
+    if (hasMatplotlib) {
+      await this.pyodide.runPythonAsync(PATCH_MATPLOTLIB);
+    }
+  }
+
+  private async patchFonts() {
+    this.pyodide.FS.mkdirTree('/usr/share/fonts/truetype');
+    const fontFiles = {
+      'STSong.ttf':
+        'https://cdn.jsdelivr.net/gh/Haixing-Hu/latex-chinese-fonts@latest/chinese/宋体/STSong.ttf',
+    };
+    for (const [filename, url] of Object.entries(fontFiles)) {
+      const buffer = await fetch(url, { cache: 'force-cache' }).then((res) => res.arrayBuffer());
+      // NOTE: 此处理论上使用 createLazyFile 更好，但 pyodide 中使用会导致报错
+      this.pyodide.FS.writeFile(`/usr/share/fonts/truetype/${filename}`, new Uint8Array(buffer));
+    }
+  }
+
+  private async isNewFile(file: File) {
+    const isSameFile = async (a: File, b: File) => {
+      // a 是传入的文件，可能使用了绝对路径或相对路径
+      // b 是解释器环境中的文件，使用绝对路径
+      if (a.name.startsWith('/')) {
+        if (a.name !== b.name) return false;
+      } else {
+        if (`/mnt/data/${a.name}` !== b.name) return false;
+      }
+
+      if (a.size !== b.size) return false;
+
+      const aBuffer = await a.arrayBuffer();
+      const bBuffer = await b.arrayBuffer();
+      const aArray = new Uint8Array(aBuffer);
+      const bArray = new Uint8Array(bBuffer);
+      const length = aArray.length;
+      for (let i = 0; i < length; i++) {
+        if (aArray[i] !== bArray[i]) return false;
+      }
+
+      return true;
+    };
+    const t = await Promise.all(this.uploadedFiles.map((f) => isSameFile(f, file)));
+    return t.every((f) => !f);
+  }
+}
+
+Comlink.expose(PythonWorker);
+
+export { PythonWorker };
+export type PythonWorkerType = typeof PythonWorker;

packages/python-interpreter/vitest.config.mts

@@ -0,0 +1,10 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    coverage: {
+      reporter: ['text', 'json', 'lcov', 'text-summary'],
+    },
+    environment: 'node',
+  },
+});