Add more types to Pyodide glue code (#2303)

* more pyodide types

* Even more types

* Fix build, address nits

* address nits, run prettier

* fix bug, add return types

src/pyodide/BUILD.bazel

@@ -164,6 +164,7 @@ wd_ts_bundle(
         "internal/topLevelEntropy/*.ts",
         "internal/topLevelEntropy/*.js",
         "types/*.ts",
+        "types/*/*.ts",
     ], allow_empty = True),
     internal_wasm_modules = ["generated/pyodide.asm.wasm"],
     schema_id = "0xbcc8f57c63814005",

src/pyodide/internal/builtin_wrappers.ts

@@ -8,7 +8,7 @@ let lastDelta = 0;
  * directories change their modification time when updated so that Python
  * doesn't use stale directory contents in its import system.
  */
-export function monotonicDateNow() {
+export function monotonicDateNow(): number {
   const now = Date.now();
   if (now === lastTime) {
     lastDelta++;
@@ -52,7 +52,7 @@ export function monotonicDateNow() {
  *      - ctypes is quite slow even by Python's standards
  *      - Normally ctypes allocates all closures up front
  */
-export function newWasmModule(buffer: Uint8Array) {
+export function newWasmModule(buffer: Uint8Array): WebAssembly.Module {
   checkCallee();
   return UnsafeEval.newWasmModule(buffer);
 }
@@ -63,7 +63,7 @@ export function newWasmModule(buffer: Uint8Array) {
  * `convertJsFunctionToWasm` in `"pyodide-internal:generated/pyodide.asm"`,
  * if it's anything else we'll bail.
  */
-function checkCallee() {
+function checkCallee(): void {
   const origPrepareStackTrace = Error.prepareStackTrace;
   let isOkay;
   try {
@@ -83,7 +83,7 @@ function checkCallee() {
  * `convertJsFunctionToWasm` or `loadModule` in `pyodide.asm.js`, `false` if not. This will set
  * the `stack` field in the error so we can read back the result there.
  */
-function prepareStackTrace(_error: Error, stack: StackItem[]) {
+function prepareStackTrace(_error: Error, stack: StackItem[]): boolean {
   // In case a logic error is ever introduced in this function, defend against
   // reentrant calls by setting `prepareStackTrace` to `undefined`.
   Error.prepareStackTrace = undefined;
@@ -105,10 +105,16 @@ function prepareStackTrace(_error: Error, stack: StackItem[]) {
   }
 }
 
-export async function wasmInstantiate(module: WebAssembly.Module | Uint8Array, imports: WebAssembly.Imports) {
-  if (!(module instanceof WebAssembly.Module)) {
+export async function wasmInstantiate(
+  mod: WebAssembly.Module | Uint8Array,
+  imports: WebAssembly.Imports,
+): Promise<{ module: WebAssembly.Module, instance: WebAssembly.Instance }> {
+  let module;
+  if (mod instanceof WebAssembly.Module) {
+    module = mod;
+  } else {
     checkCallee();
-    module = UnsafeEval.newWasmModule(module);
+    module = UnsafeEval.newWasmModule(mod);
   }
   const instance = new WebAssembly.Instance(module, imports);
   return { module, instance };

src/pyodide/internal/jaeger.ts

@@ -1,10 +1,9 @@
-// @ts-nocheck
 import { default as internalJaeger } from "pyodide-internal:internalJaeger";
 
 /**
  * Used for tracing via Jaeger.
  */
-export function enterJaegerSpan(span, callback) {
+export function enterJaegerSpan<T>(span: string, callback: () => T): T {
   if (!internalJaeger.traceId) {
     // Jaeger tracing not enabled or traceId is not present in request.
     return callback();

src/pyodide/internal/loadPackage.ts

@@ -1,4 +1,3 @@
-// @ts-nocheck
 /**
  * This file contains code that roughly replaces pyodide.loadPackage, with workerd-specific
  * optimizations:
@@ -12,16 +11,29 @@
 
 import { default as LOCKFILE } from "pyodide-internal:generated/pyodide-lock.json";
 import { WORKERD_INDEX_URL } from "pyodide-internal:metadata";
-import { SITE_PACKAGES, LOAD_WHEELS_FROM_R2, getSitePackagesPath } from "pyodide-internal:setupPackages";
+import {
+  SITE_PACKAGES,
+  LOAD_WHEELS_FROM_R2,
+  getSitePackagesPath,
+} from "pyodide-internal:setupPackages";
 import { parseTarInfo } from "pyodide-internal:tar";
 import { default as DiskCache } from "pyodide-internal:disk_cache";
 import { createTarFS } from "pyodide-internal:tarfs";
 
-async function decompressArrayBuffer(arrBuf) {
-  return await new Response(new Response(arrBuf).body.pipeThrough(new DecompressionStream("gzip"))).arrayBuffer();
+async function decompressArrayBuffer(
+  arrBuf: ArrayBuffer,
+): Promise<ArrayBuffer> {
+  const resp = new Response(arrBuf);
+  if (resp && resp.body) {
+    return await new Response(
+      resp.body.pipeThrough(new DecompressionStream("gzip")),
+    ).arrayBuffer();
+  } else {
+    throw new Error("Failed to decompress array buffer");
+  }
 }
 
-async function loadBundle(requirement) {
+async function loadBundle(requirement: string): Promise<[string, ArrayBuffer]> {
   // first check if the disk cache has what we want
   const filename = LOCKFILE["packages"][requirement]["file_name"];
   const cached = DiskCache.get(filename);
@@ -39,18 +51,15 @@ async function loadBundle(requirement) {
 
   DiskCache.put(filename, compressed);
   return [requirement, decompressed];
-};
+}
 
 /**
  * ArrayBufferReader wraps around an arrayBuffer in a way that tar.js is able to read from
  */
 class ArrayBufferReader {
-  constructor(arrayBuffer) {
-    this.arrayBuffer = arrayBuffer;
-  }
+  constructor(private arrayBuffer: ArrayBuffer) {}
 
-  read(offset, buf){
-    // buf is a Uint8Array
+  read(offset: number, buf: Uint8Array): number {
     const size = this.arrayBuffer.byteLength;
     if (offset >= size || offset < 0) {
       return 0;
@@ -64,7 +73,7 @@ class ArrayBufferReader {
   }
 }
 
-export async function loadPackages(Module, requirements) {
+export async function loadPackages(Module: Module, requirements: string[]) {
   if (!LOAD_WHEELS_FROM_R2) return;
 
   let loadPromises = [];

src/pyodide/internal/metadata.ts

@@ -1,4 +1,3 @@
-// @ts-nocheck
 import { default as MetadataReader } from "pyodide-internal:runtime-generated/metadata";
 export { default as LOCKFILE } from "pyodide-internal:generated/pyodide-lock.json";
 import { default as PYODIDE_BUCKET } from "pyodide-internal:generated/pyodide-bucket.json";

src/pyodide/internal/setupPackages.ts

@@ -1,4 +1,3 @@
-// @ts-nocheck
 import { parseTarInfo } from "pyodide-internal:tar";
 import { createTarFS } from "pyodide-internal:tarfs";
 import { createMetadataFS } from "pyodide-internal:metadatafs";
@@ -14,12 +13,12 @@ const canonicalizeNameRegex = /[-_.]+/g;
  * @returns The canonicalize package name.
  * @private
  */
-function canonicalizePackageName(name) {
+function canonicalizePackageName(name: string): string {
   return name.replace(canonicalizeNameRegex, "-").toLowerCase();
 }
 
 // The "name" field in the lockfile is not canonicalized
-const STDLIB_PACKAGES = Object.values(LOCKFILE.packages)
+const STDLIB_PACKAGES: string[] = Object.values(LOCKFILE.packages)
   .filter(({ install_dir }) => install_dir === "stdlib")
   .map(({ name }) => canonicalizePackageName(name));
 
@@ -28,6 +27,9 @@ const STDLIB_PACKAGES = Object.values(LOCKFILE.packages)
  * directory generated for each worker.
  */
 class SitePackagesDir {
+  public rootInfo: FSInfo;
+  public soFiles: string[][];
+  public loadedRequirements: Set<string>;
   constructor() {
     this.rootInfo = {
       children: new Map(),
@@ -49,14 +51,14 @@ class SitePackagesDir {
    * If a file or directory already exists, an error is thrown.
    * @param {TarInfo} overlayInfo The directory that is to be "copied" into site-packages
    */
-  mountOverlay(overlayInfo) {
-    overlayInfo.children.forEach((val, key) => {
-      if (this.rootInfo.children.has(key)) {
+  mountOverlay(overlayInfo: FSInfo): void {
+    overlayInfo.children!.forEach((val, key) => {
+      if (this.rootInfo.children!.has(key)) {
         throw new Error(
           `File/folder ${key} being written by multiple packages`,
         );
       }
-      this.rootInfo.children.set(key, val);
+      this.rootInfo.children!.set(key, val);
     });
   }
 
@@ -67,7 +69,11 @@ class SitePackagesDir {
    * @param {List<String>} soFiles A list of .so files contained in the small bundle
    * @param {String} requirement The canonicalized package name this small bundle corresponds to
    */
-  addSmallBundle(tarInfo, soFiles, requirement) {
+  addSmallBundle(
+    tarInfo: FSInfo,
+    soFiles: string[],
+    requirement: string,
+  ): void {
     for (const soFile of soFiles) {
       this.soFiles.push(soFile.split("/"));
     }
@@ -82,7 +88,11 @@ class SitePackagesDir {
    * @param {List<String>} soFiles A list of .so files contained in the big bundle
    * @param {List<String>} requirements canonicalized list of packages to pick from the big bundle
    */
-  addBigBundle(tarInfo, soFiles, requirements) {
+  addBigBundle(
+    tarInfo: FSInfo,
+    soFiles: string[],
+    requirements: Set<string>,
+  ): void {
     // add all the .so files we will need to preload from the big bundle
     for (const soFile of soFiles) {
       // If folder is in list of requirements include .so file in list to preload.
@@ -93,15 +103,15 @@ class SitePackagesDir {
     }
 
     for (const req of requirements) {
-      const child = tarInfo.children.get(req);
+      const child = tarInfo.children!.get(req);
       if (!child) {
         throw new Error(`Requirement ${req} not found in pyodide packages tar`);
       }
       this.mountOverlay(child);
       this.loadedRequirements.add(req);
     }
   }
-};
+}
 
 /**
  * This stitches together the view of the site packages directory. Each
@@ -112,14 +122,16 @@ class SitePackagesDir {
  * This also returns the list of soFiles in the resulting site-packages
  * directory so we can preload them.
  */
-export function buildSitePackages(requirements) {
+export function buildSitePackages(
+  requirements: Set<string>,
+): [SitePackagesDir, boolean] {
   const [bigTarInfo, bigTarSoFiles] = parseTarInfo();
 
   let LOAD_WHEELS_FROM_R2 = true;
   let requirementsInBigBundle = new Set([...STDLIB_PACKAGES]);
-  if (bigTarInfo.children.size > 10) {
+  if (bigTarInfo.children!.size > 10) {
     LOAD_WHEELS_FROM_R2 = false;
-    requirements.forEach(r => requirementsInBigBundle.add(r));
+    requirements.forEach((r) => requirementsInBigBundle.add(r));
   }
 
   const res = new SitePackagesDir();
@@ -135,26 +147,28 @@ export function buildSitePackages(requirements) {
  *
  * TODO: stop using loadPackage in workerd.
  */
-export function patchLoadPackage(pyodide) {
+export function patchLoadPackage(pyodide: { loadPackage: Function }): void {
   pyodide.loadPackage = disabledLoadPackage;
   return;
 }
 
 function disabledLoadPackage() {
-  throw new Error("pyodide.loadPackage is disabled because packages are encoded in the binary");
+  throw new Error(
+    "pyodide.loadPackage is disabled because packages are encoded in the binary",
+  );
 }
 
 /**
  * Get the set of transitive requirements from the REQUIREMENTS metadata.
  */
-function getTransitiveRequirements() {
+function getTransitiveRequirements(): Set<string> {
   const requirements = REQUIREMENTS.map(canonicalizePackageName);
   // resolve transitive dependencies of requirements and if IN_WORKERD install them from the cdn.
   const packageDatas = recursiveDependencies(LOCKFILE, requirements);
   return new Set(packageDatas.map(({ name }) => canonicalizePackageName(name)));
 }
 
-export function getSitePackagesPath(Module) {
+export function getSitePackagesPath(Module: Module): string {
   const pymajor = Module._py_version_major();
   const pyminor = Module._py_version_minor();
   return `/session/lib/python${pymajor}.${pyminor}/site-packages`;
@@ -168,7 +182,7 @@ export function getSitePackagesPath(Module) {
  * details, so even though we want these directories to be on sys.path, we
  * handle that separately in adjustSysPath.
  */
-export function mountLib(Module, info) {
+export function mountLib(Module: Module, info: FSInfo): void {
   const tarFS = createTarFS(Module);
   const mdFS = createMetadataFS(Module);
   const site_packages = getSitePackagesPath(Module);
@@ -187,15 +201,18 @@ export function mountLib(Module, info) {
  * Add the directories created by mountLib to sys.path.
  * Has to run after the runtime is initialized.
  */
-export function adjustSysPath(Module) {
+export function adjustSysPath(Module: Module): void {
   const site_packages = getSitePackagesPath(Module);
   simpleRunPython(
     Module,
     `import sys; sys.path.append("/session/metadata"); sys.path.append("${site_packages}"); del sys`,
   );
 }
 
-function recursiveDependencies(lockfile, names) {
+function recursiveDependencies(
+  lockfile: PackageLock,
+  names: string[],
+): PackageDeclaration[] {
   const toLoad = new Map();
   for (const name of names) {
     addPackageToLoad(lockfile, name, toLoad);
@@ -210,7 +227,11 @@ function recursiveDependencies(lockfile, names) {
  * @param toLoad The set of names of packages to load
  * @private
  */
-function addPackageToLoad(lockfile, name, toLoad) {
+function addPackageToLoad(
+  lockfile: PackageLock,
+  name: string,
+  toLoad: Map<string, PackageDeclaration>,
+): void {
   const normalizedName = canonicalizePackageName(name);
   if (toLoad.has(normalizedName)) {
     return;
@@ -219,7 +240,7 @@ function addPackageToLoad(lockfile, name, toLoad) {
   if (!pkgInfo) {
     throw new Error(
       `It appears that a package ("${name}") you requested is not available yet in workerd. \n` +
-        "If you would like this package to be included, please open an issue at https://github.com/cloudflare/workerd/discussions/new?category=python-packages.",
+      "If you would like this package to be included, please open an issue at https://github.com/cloudflare/workerd/discussions/new?category=python-packages.",
     );
   }
 
@@ -232,4 +253,6 @@ function addPackageToLoad(lockfile, name, toLoad) {
 
 export { REQUIREMENTS };
 export const TRANSITIVE_REQUIREMENTS = getTransitiveRequirements();
-export const [SITE_PACKAGES, LOAD_WHEELS_FROM_R2] = buildSitePackages(TRANSITIVE_REQUIREMENTS);
+export const [SITE_PACKAGES, LOAD_WHEELS_FROM_R2] = buildSitePackages(
+  TRANSITIVE_REQUIREMENTS,
+);

src/pyodide/internal/tar.ts

@@ -45,7 +45,7 @@ function decodeHeader(buf, reader) {
   };
 }
 
-export function parseTarInfo(reader = TarReader) {
+export function parseTarInfo(reader = TarReader): [FSInfo, string[]] {
   const directories = [];
   const soFiles = [];
   const root = {

src/pyodide/internal/tarfs.ts

@@ -24,7 +24,10 @@ const FSOps = {
     return parent.info.children.get(name);
   },
   read(stream, position, buffer) {
-    return stream.node.info.reader.read(stream.node.contentsOffset + position, buffer);
+    return stream.node.info.reader.read(
+      stream.node.contentsOffset + position,
+      buffer,
+    );
   },
 };
 

src/pyodide/internal/topLevelEntropy/entropy_import_context.py.d.ts

@@ -1,2 +1,2 @@
-declare const buf: ArrayBuffer
+declare const buf: ArrayBuffer;
 export default buf;

src/pyodide/internal/topLevelEntropy/entropy_patches.py.d.ts

@@ -1,2 +1,2 @@
-declare const buf: ArrayBuffer
+declare const buf: ArrayBuffer;
 export default buf;

src/pyodide/internal/topLevelEntropy/import_patch_manager.py.d.ts

@@ -1,2 +1,2 @@
-declare const buf: ArrayBuffer
+declare const buf: ArrayBuffer;
 export default buf;