new/frontends/react/src/PyodideBackend.ts

const PYODIDE_URL = "https://cdn.jsdelivr.net/pyodide/v0.26.1/full/pyodide.mjs";

import managerCode from "../../../backend/src/manager.py?raw";
// import logicCode from "../../../backend/src/logic.py?raw";


export class PyodideBackend {
  private initialized: boolean = false;
  private initPromise: Promise<void> | null = null;

  private chain: Promise<void> = Promise.resolve();

  private pyodide: any = null;
  private manager: any = null;

  async init(): Promise<void> {
    if (this.initialized) {
      return;
    }
    if (this.initPromise) {
      return this.initPromise;
    }

    this.initPromise = (async () => {
      const { loadPyodide } = await import(/* @vite-ignore */ PYODIDE_URL);

      this.pyodide = await loadPyodide({
        indexURL: "https://cdn.jsdelivr.net/pyodide/v0.26.1/full/"
      });

      await this.pyodide.loadPackage(["numpy", "micropip", "sqlite3"]);
      await this.pyodide.runPythonAsync("import micropip; await micropip.install('tinygrad')");

      this.pyodide.FS.writeFile("manager.py", managerCode);
      this.pyodide.runPython(`from manager import Manager; manager = Manager();`);
      this.manager = this.pyodide.globals.get("manager");

      if (!this.manager) {
        throw new Error("Failed to initialize pyodide manager");
      }

      console.log("Pyodide initialized");
      this.initialized = true;
    })();

    return this.initPromise;
  }

  async handleImageFolderUpload(files: File[]): Promise<string | null> {
    if (files.length === 0) {
      return null;
    }

    const baseDir = "/dataset";
    this.pyodide.FS.mkdirTree(baseDir);

    let rootDir: string | null = null;

    console.log("Uploading...");
    for (const f of files) {
      const relPath = (f as any).webkitRelativePath;

      if (!relPath) {
        throw new Error("All files must come from a directory upload");
      }

      const parts = relPath.split("/");

      if (parts.length < 2) {
        throw new Error(`Invalid relative path: ${relPath}`);
      }

      const currentRoot = parts[0];

      if (rootDir === null) {
        rootDir = currentRoot;
      } else if (currentRoot !== rootDir) {
        throw new Error(
          `Multiple root directories detected: '${rootDir}' and '${currentRoot}'`
        );
      }

      // strip the original root directory
      const strippedPath = parts.slice(1).join("/");
      const fullPath = `${baseDir}/${strippedPath}`;

      const dir = fullPath.split("/").slice(0, -1).join("/");
      this.pyodide.FS.mkdirTree(dir);

      const buf = await f.arrayBuffer();
      const u8 = new Uint8Array(buf);

      this.pyodide.FS.writeFile(fullPath, u8);
    }

    await this.handleCall<void>("handle_index_imagefolder");

    return rootDir;
  }

  async handleGetRandomSample(): Promise<{ imageUrl: string; label: string }> {
    const { imagePath, label } = await this.handleCall<{ imagePath: string; label: string }>(
      "handle_get_random_sample"
    );

    const data = this.pyodide.FS.readFile(imagePath);

    const mime = this.getImageMimeType(imagePath);
    const blob = new Blob([data], { type: mime });

    const url = URL.createObjectURL(blob);

    return { imageUrl: url, label };
  }

  private getImageMimeType(filename: string): string {
    const ext = this.getFileExtension(filename);

    if (ext === "png") {
      return "image/png";
    }

    if (ext === "jpg" || ext === "jpeg") {
      return "image/jpeg";
    }

    if (ext === "webp") {
      return "image/webp";
    }

    if (ext === "gif") {
      return "image/gif";
    }

    return "application/octet-stream";
  }

  private getFileExtension(filename: string): string {
    const parts = filename.split('.');
    if (parts.length < 2) {
      return '';
    }
    return parts[parts.length - 1].toLowerCase();
  }

  private enqueue<T>(fn: () => Promise<T>): Promise<T> {
    const next = this.chain.then(() => fn(), () => fn());
    this.chain = next.then(() => undefined, () => undefined);
    return next;
  }

  private async handleCall<T>(methodName: string, args?: unknown): Promise<T> {
    await this.init();

    return this.enqueue(async () => {
      let pyArgs: any = null;
      let pyResult: any = null;

      try {
        if (args !== undefined) {
          pyArgs = this.pyodide.toPy(args);
        }

        const fn = this.manager[methodName];
        if (!fn) {
          throw new Error(`Manager has no method named ${methodName}`);
        }

        pyResult = args === undefined ? fn.call(this.manager) : fn.call(this.manager, pyArgs);

        const result = pyResult && typeof pyResult.toJs === "function" 
          ? pyResult.toJs({ dict_converter: Object.fromEntries })
          : pyResult;
        
        return result as T;

      } finally {
        if (pyArgs && typeof pyArgs.destroy === "function") {
          pyArgs.destroy();
        }
        if (pyResult && typeof pyResult.destroy === "function") {
          pyResult.destroy();
        }
      }
    });
  }
};