Repository: pyodide/pyodide
Pull request #5481: MAINT Add context manager for TS and use it in loadPackage
Author: login: ryanking13
Created: 2025-02-27T14:54:20Z | Merged: 2025-03-04T06:14:44Z

### Description

This PR implements a Python-like context manager in Typescript, and apply it in the `loadPackage`.

The high-level function `withContext` can be used to create a context and can be used such as:

```typescript
const result = withContext(
    doSetup,
    doCleanup,
    doRealwork,
);
```

This context manager is applied to `loadPackage` to set `messageCallback` and `errorCallback` as a context. Previously we needed drill this parameters into child functions, but not a top-level context manager handles it.

There's no functional change for most people, but this makes loadPackage a little more thread-unsafe. But we never guaranteed thread safety for loadPackage, so I think it is not a problem.

### Checklists

- [x] Add / update tests



Diff:
diff --git a/src/js/common/contextManager.ts b/src/js/common/contextManager.ts
new file mode 100644
index 00000000000..e926237a9d3
--- /dev/null
+++ b/src/js/common/contextManager.ts
@@ -0,0 +1,53 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => void,
+  cleanup: () => void,
+  callback: () => T,
+): T {
+  setup();
+  let result: T;
+  try {
+    result = callback();
+  } catch (e) {
+    cleanup();
+    throw e;
+  }
+  if (result instanceof Promise) {
+    return result.finally(() => cleanup()) as T;
+  }
+
+  cleanup();
+  return result;
+}
+
+// A type for a function that wraps another function and preserves its type
+type FunctionWrapper = <T extends (...args: any[]) => any>(
+  fn: T,
+) => (...args: Parameters<T>) => ReturnType<T>;
+
+/**
+ * Creates a function wrapper that sets up a context before calling the function
+ * and cleans up afterwards. Works with both synchronous and asynchronous functions.
+ *
+ * @param setup Function to call before the wrapped function
+ * @param cleanup Function to call after the wrapped function
+ * @returns A function that wraps another function with the context
+ */
+export function createContextWrapper(
+  setup: () => void,
+  cleanup: () => void,
+): FunctionWrapper {
+  return function (fn) {
+    return function (this: any, ...args: Parameters<typeof fn>) {
+      return withContext(setup, cleanup, () => fn.apply(this, args));
+    };
+  };
+}
diff --git a/src/js/load-package.ts b/src/js/load-package.ts
index 564fc580578..7bc532ef92f 100644
--- a/src/js/load-package.ts
+++ b/src/js/load-package.ts
@@ -25,6 +25,7 @@ import {
   ensureDirNode,
 } from "./compat";
 import { Installer } from "./installer";
+import { createContextWrapper } from "./common/contextManager";
 
 /**
  * Initialize the packages index. This is called as early as possible in
@@ -182,12 +183,29 @@ export class PackageManager {
     } = {
       checkIntegrity: true,
     },
+  ): Promise<PackageData[]> {
+    const wrappedLoadPackage = this.setCallbacks(
+      options.messageCallback,
+      options.errorCallback,
+    )(this.loadPackageInner);
+
+    return wrappedLoadPackage.call(this, names, options);
+  }
+
+  public async loadPackageInner(
+    names: string | PyProxy | string[],
+    options: {
+      messageCallback?: (message: string) => void;
+      errorCallback?: (message: string) => void;
+      checkIntegrity?: boolean;
+    } = {
+      checkIntegrity: true,
+    },
   ): Promise<Array<PackageData>> {
     const loadedPackageData = new Set<InternalPackageData>();
-    const { messageCallback, errorCallback } = options;
     const pkgNames = toStringArray(names);
 
-    const toLoad = this.recursiveDependencies(pkgNames, errorCallback);
+    const toLoad = this.recursiveDependencies(pkgNames);
 
     for (const [_, { name, normalizedName, channel }] of toLoad) {
       const loadedChannel = this.getLoadedPackageChannel(name);
@@ -197,22 +215,18 @@ export class PackageManager {
       // If uri is from the default channel, we assume it was added as a
       // dependency, which was previously overridden.
       if (loadedChannel === channel || channel === this.defaultChannel) {
-        this.logStdout(
-          `${name} already loaded from ${loadedChannel}`,
-          messageCallback,
-        );
+        this.logStdout(`${name} already loaded from ${loadedChannel}`);
       } else {
         this.logStderr(
           `URI mismatch, attempting to load package ${name} from ${channel} ` +
             `while it is already loaded from ${loadedChannel}. To override a dependency, ` +
             `load the custom package first.`,
-          errorCallback,
         );
       }
     }
 
     if (toLoad.size === 0) {
-      this.logStdout("No new packages to load", messageCallback);
+      this.logStdout("No new packages to load");
       return [];
     }
 
@@ -222,7 +236,7 @@ export class PackageManager {
     const failed = new Map<string, Error>();
     const releaseLock = await this._lock();
     try {
-      this.logStdout(`Loading ${packageNames}`, messageCallback);
+      this.logStdout(`Loading ${packageNames}`);
       for (const [_, pkg] of toLoad) {
         if (this.getLoadedPackageChannel(pkg.name)) {
           // Handle the race condition where the package was loaded between when
@@ -253,18 +267,15 @@ export class PackageManager {
         const successNames = Array.from(loadedPackageData, (pkg) => pkg.name)
           .sort()
           .join(", ");
-        this.logStdout(`Loaded ${successNames}`, messageCallback);
+        this.logStdout(`Loaded ${successNames}`);
       }
 
       if (failed.size > 0) {
         const failedNames = Array.from(failed.keys()).sort().join(", ");
-        this.logStdout(`Failed to load ${failedNames}`, messageCallback);
+        this.logStdout(`Failed to load ${failedNames}`);
         for (const [name, err] of failed) {
-          this.logStderr(
-            `The following error occurred while loading ${name}:`,
-            errorCallback,
-          );
-          this.logStderr(err.message, errorCallback);
+          this.logStderr(`The following error occurred while loading ${name}:`);
+          this.logStderr(err.message);
         }
       }
 
@@ -325,9 +336,8 @@ export class PackageManager {
    * @returns The map of package names to PackageLoadMetadata
    * @private
    */
-  public recursiveDependencies(
+  private recursiveDependencies(
     names: string[],
-    errorCallback?: (err: string) => void,
   ): Map<string, PackageLoadMetadata> {
     const toLoad: Map<string, PackageLoadMetadata> = new Map();
     for (let name of names) {
@@ -345,7 +355,6 @@ export class PackageManager {
           `Loading same package ${pkgname} from ${channel} and ${
             toLoad.get(pkgname)!.channel
           }`,
-          errorCallback,
         );
         continue;
       }
@@ -534,12 +543,31 @@ export class PackageManager {
     return channel;
   }
 
-  public logStdout(message: string, logger?: (message: string) => void) {
-    logger ? logger(message) : this.stdout(message);
+  public setCallbacks(
+    stdout?: (message: string) => void,
+    stderr?: (message: string) => void,
+  ) {
+    const originalStdout = this.stdout;
+    const originalStderr = this.stderr;
+
+    return createContextWrapper(
+      () => {
+        this.stdout = stdout || originalStdout;
+        this.stderr = stderr || originalStderr;
+      },
+      () => {
+        this.stdout = originalStdout;
+        this.stderr = originalStderr;
+      },
+    );
+  }
+
+  public logStdout(message: string) {
+    this.stdout(message);
   }
 
-  public logStderr(message: string, logger?: (message: string) => void) {
-    logger ? logger(message) : this.stderr(message);
+  public logStderr(message: string) {
+    this.stderr(message);
   }
 }
 
@@ -599,8 +627,6 @@ if (typeof API !== "undefined" && typeof Module !== "undefined") {
   loadedPackages = singletonPackageManager.loadedPackages;
 
   // TODO: Find a better way to register these functions
-  API.recursiveDependencies =
-    singletonPackageManager.recursiveDependencies.bind(singletonPackageManager);
   API.setCdnUrl = singletonPackageManager.setCdnUrl.bind(
     singletonPackageManager,
   );
diff --git a/src/js/test/unit/common/contextManager.test.ts b/src/js/test/unit/common/contextManager.test.ts
new file mode 100644
index 00000000000..0b48eec57b3
--- /dev/null
+++ b/src/js/test/unit/common/contextManager.test.ts
@@ -0,0 +1,240 @@
+import * as chai from "chai";
+import chaiAsPromised from "chai-as-promised";
+
+chai.use(chaiAsPromised);
+const assert = chai.assert;
+const expect = chai.expect;
+
+import {
+  withContext,
+  createContextWrapper,
+} from "../../../common/contextManager";
+
+describe("withContext", () => {
+  it("synchronous execution", () => {
+    // Setup tracking variables
+    let setupCalled = false;
+    let cleanupCalled = false;
+    let callbackCalled = false;
+
+    const result = withContext(
+      () => {
+        setupCalled = true;
+      },
+      () => {
+        cleanupCalled = true;
+      },
+      () => {
+        callbackCalled = true;
+        return "test result";
+      },
+    );
+
+    assert.isTrue(setupCalled);
+    assert.isTrue(cleanupCalled);
+    assert.isTrue(callbackCalled);
+    assert.equal(result, "test result");
+  });
+
+  it("asynchronous execution", async () => {
+    // Setup tracking variables
+    let setupCalled = false;
+    let cleanupCalled = false;
+    let callbackCalled = false;
+
+    const result = await withContext(
+      () => {
+        setupCalled = true;
+      },
+      () => {
+        cleanupCalled = true;
+      },
+      async () => {
+        callbackCalled = true;
+        await new Promise((resolve) => setTimeout(resolve, 10));
+        return "async result";
+      },
+    );
+
+    assert.isTrue(setupCalled);
+    assert.isTrue(cleanupCalled);
+    assert.isTrue(callbackCalled);
+    assert.equal(result, "async result");
+  });
+
+  it("cleanup is called even when callback throws", () => {
+    // Setup tracking variables
+    let setupCalled = false;
+    let cleanupCalled = false;
+
+    assert.throws(
+      () =>
+        withContext(
+          () => {
+            setupCalled = true;
+          },
+          () => {
+            cleanupCalled = true;
+          },
+          () => {
+            throw new Error("Test error");
+          },
+        ),
+      "Test error",
+    );
+
+    assert.isTrue(setupCalled);
+    assert.isTrue(cleanupCalled);
+  });
+
+  it("cleanup is called even when async callback rejects", async () => {
+    // Setup tracking variables
+    let setupCalled = false;
+    let cleanupCalled = false;
+
+    const promise = withContext(
+      () => {
+        setupCalled = true;
+      },
+      () => {
+        cleanupCalled = true;
+      },
+      async () => {
+        await new Promise((resolve) => setTimeout(resolve, 10));
+        throw new Error("Async test error");
+      },
+    );
+
+    await expect(promise).to.be.rejectedWith("Async test error");
+    assert.isTrue(setupCalled);
+    assert.isTrue(cleanupCalled);
+  });
+
+  it("raise on setup error", () => {
+    assert.throws(
+      () =>
+        withContext(
+          () => {
+            throw new Error("Setup error");
+          },
+          () => {},
+          () => {},
+        ),
+      "Setup error",
+    );
+  });
+
+  it("raise on cleanup error", () => {
+    assert.throws(
+      () =>
+        withContext(
+          () => {},
+          () => {
+            throw new Error("Cleanup error");
+          },
+          () => {},
+        ),
+      "Cleanup error",
+    );
+  });
+});
+
+describe("createContextWrapper", () => {
+  it("wrapper for sync function", () => {
+    // Setup tracking variables
+    let setupCalled = false;
+    let cleanupCalled = false;
+
+    function testFn(x: number, y: number) {
+      return x + y;
+    }
+
+    const wrappedFn = createContextWrapper(
+      () => {
+        setupCalled = true;
+      },
+      () => {
+        cleanupCalled = true;
+      },
+    )(testFn);
+
+    const result = wrappedFn(5, 3);
+
+    assert.isTrue(setupCalled);
+    assert.isTrue(cleanupCalled);
+    assert.equal(result, 8);
+  });
+
+  it("wrapper for async function", async () => {
+    // Setup tracking variables
+    let setupCalled = false;
+    let cleanupCalled = false;
+
+    async function testAsyncFn(x: number, y: number) {
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      return x * y;
+    }
+
+    const wrappedFn = createContextWrapper(
+      () => {
+        setupCalled = true;
+      },
+      () => {
+        cleanupCalled = true;
+      },
+    )(testAsyncFn);
+
+    const result = await wrappedFn(4, 7);
+
+    assert.isTrue(setupCalled);
+    assert.isTrue(cleanupCalled);
+    assert.equal(result, 28);
+  });
+
+  it("cleanup is called when wrapped function throws", () => {
+    // Setup tracking variables
+    let setupCalled = false;
+    let cleanupCalled = false;
+
+    function throwingFn() {
+      throw new Error("Function error");
+    }
+
+    const wrappedFn = createContextWrapper(
+      () => {
+        setupCalled = true;
+      },
+      () => {
+        cleanupCalled = true;
+      },
+    )(throwingFn);
+
+    assert.throws(() => wrappedFn(), "Function error");
+    assert.isTrue(setupCalled);
+    assert.isTrue(cleanupCalled);
+  });
+
+  it("cleanup is called when async wrapped function rejects", async () => {
+    // Setup tracking variables
+    let setupCalled = false;
+    let cleanupCalled = false;
+
+    async function throwingAsyncFn() {
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      throw new Error("Async function error");
+    }
+
+    const wrappedFn = createContextWrapper(
+      () => {
+        setupCalled = true;
+      },
+      () => {
+        cleanupCalled = true;
+      },
+    )(throwingAsyncFn);
+
+    await assert.isRejected(wrappedFn(), "Async function error");
+    assert.isTrue(setupCalled);
+    assert.isTrue(cleanupCalled);
+  });
+});
diff --git a/src/js/test/unit/package-manager.test.ts b/src/js/test/unit/package-manager.test.ts
index 8ef36b8dd26..462cb8fa37b 100644
--- a/src/js/test/unit/package-manager.test.ts
+++ b/src/js/test/unit/package-manager.test.ts
@@ -33,7 +33,7 @@ describe("logStdout and logStderr", () => {
     errorSpy.restore();
   });
 
-  it("Should use the provided logger if one is provided", () => {
+  it("Should be overwritten when setCallbacks is called", () => {
     const mockApi = genMockAPI();
     const mockMod = genMockModule();
 
@@ -42,8 +42,13 @@ describe("logStdout and logStderr", () => {
     const stdoutLogger = sinon.spy();
     const stderrLogger = sinon.spy();
 
-    pm.logStdout("stdout", stdoutLogger);
-    pm.logStderr("stderr", stderrLogger);
+    pm.setCallbacks(
+      stdoutLogger,
+      stderrLogger,
+    )(() => {
+      pm.logStdout("stdout");
+      pm.logStderr("stderr");
+    })();
 
     chai.assert.isTrue(stdoutLogger.calledOnce);
     chai.assert.isTrue(stderrLogger.calledOnce);
diff --git a/src/js/types.ts b/src/js/types.ts
index 4af2378ad87..c258ecd464b 100644
--- a/src/js/types.ts
+++ b/src/js/types.ts
@@ -492,10 +492,6 @@ export interface API {
     installDir: string,
     metadata?: ReadonlyMap<string, string>,
   ) => Promise<void>;
-  recursiveDependencies: (
-    names: string[],
-    errorCallback: (err: string) => void,
-  ) => Map<string, PackageLoadMetadata>;
   _Comlink: any;
 
   dsodir: string;


Reviews:
user: login: ryanking13
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-02-27T14:56:32Z
user: login: hoodmane
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-02-27T18:02:14Z
user: login: hoodmane
body: Looks good to me aside from a few comments.
state: APPROVED
author_association: MEMBER
submitted_at: 2025-02-27T20:51:59Z
user: login: ryanking13
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-02T01:44:21Z
user: login: ryanking13
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-02T02:01:13Z
user: login: ryanking13
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-02T02:01:38Z
user: login: hoodmane
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-02T08:19:40Z
user: login: ryanking13
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-02T10:46:37Z
user: login: hoodmane
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-02T10:49:28Z
user: login: hoodmane
body: Let's maybe go though some more rounds of nits on the typescript types but generally looks good to me.
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-02T10:56:53Z
user: login: hoodmane
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-02T11:17:28Z
user: login: ryanking13
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-02T11:38:20Z

Review comments:
diff_hunk: @@ -599,8 +626,6 @@ if (typeof API !== "undefined" && typeof Module !== "undefined") {
   loadedPackages = singletonPackageManager.loadedPackages;
 
   // TODO: Find a better way to register these functions
-  API.recursiveDependencies =
path: src/js/load-package.ts
user: login: ryanking13
body: @hoodmane I don't remember why we've exposed this API (https://github.com/pyodide/pyodide/pull/4555). Are you using it somewhere downstream?
created_at: 2025-02-27T14:56:31Z
author_association: MEMBER
diff_hunk: @@ -599,8 +626,6 @@ if (typeof API !== "undefined" && typeof Module !== "undefined") {
   loadedPackages = singletonPackageManager.loadedPackages;
 
   // TODO: Find a better way to register these functions
-  API.recursiveDependencies =
path: src/js/load-package.ts
user: login: hoodmane
body: Doesn't look like it. I didn't explain what I was up to in #4555 and you didn't ask so I guess we'll never know...
created_at: 2025-02-27T18:02:13Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,56 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => any,
+  cleanup: () => void,
+  callback: () => T | Promise<T>,
+): T | Promise<T> {
path: src/js/common/contextManager.ts
user: login: hoodmane
body: I think the following is better:
```suggestion
export function withContext<T>(
  setup: () => any,
  cleanup: () => void,
  callback: () => T ,
): T {
```

created_at: 2025-02-27T19:41:59Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,56 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => any,
+  cleanup: () => void,
+  callback: () => T | Promise<T>,
+): T | Promise<T> {
+  setup();
+  try {
+    const result = callback();
+    if (result instanceof Promise) {
+      return result.finally(() => cleanup());
+    } else {
+      cleanup();
+      return result;
+    }
+  } catch (e) {
+    cleanup();
+    throw e;
+  }
path: src/js/common/contextManager.ts
user: login: hoodmane
body: This is a bit better I think:
```suggestion
  let result;
  try {
    result = callback();
  } catch (e) {
    cleanup();
    throw e;
  }
  if (result instanceof Promise) {
    return result.finally(() => cleanup());
  }
  cleanup();
  return result;

```

created_at: 2025-02-27T19:45:31Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,56 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => any,
+  cleanup: () => void,
+  callback: () => T | Promise<T>,
+): T | Promise<T> {
+  setup();
+  try {
+    const result = callback();
+    if (result instanceof Promise) {
+      return result.finally(() => cleanup());
+    } else {
+      cleanup();
+      return result;
+    }
+  } catch (e) {
+    cleanup();
+    throw e;
+  }
+}
+
+/**
+ * Creates a function wrapper that sets up a context before calling the function
+ * and cleans up afterwards. Works with both synchronous and asynchronous functions.
+ *
+ * @param setup Function to call before the wrapped function
+ * @param cleanup Function to call after the wrapped function
+ * @returns A function that wraps another function with the context
+ */
+export function createContextWrapper(setup: () => any, cleanup: () => void) {
+  return <T extends (...args: any[]) => any>(fn: T) => {
path: src/js/common/contextManager.ts
user: login: hoodmane
body: I think the following should work:
```suggestion
  return withContext.bind(null, setup, cleanup);
```

created_at: 2025-02-27T20:49:31Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,212 @@
+import * as chai from "chai";
+import chaiAsPromised from "chai-as-promised";
+
+chai.use(chaiAsPromised);
+const assert = chai.assert;
+const expect = chai.expect;
+
+import {
+  withContext,
+  createContextWrapper,
+} from "../../../common/contextManager";
+
+describe("withContext", () => {
+  it("synchronous execution", () => {
path: src/js/test/unit/common/contextManager.test.ts
user: login: hoodmane
body: Can you add tests for when setup and cleanup raise? Synchronously and asynchronously?
created_at: 2025-02-27T20:51:41Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,56 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => any,
+  cleanup: () => void,
+  callback: () => T | Promise<T>,
+): T | Promise<T> {
+  setup();
+  try {
+    const result = callback();
+    if (result instanceof Promise) {
+      return result.finally(() => cleanup());
+    } else {
+      cleanup();
+      return result;
+    }
+  } catch (e) {
+    cleanup();
+    throw e;
+  }
+}
+
+/**
+ * Creates a function wrapper that sets up a context before calling the function
+ * and cleans up afterwards. Works with both synchronous and asynchronous functions.
+ *
+ * @param setup Function to call before the wrapped function
+ * @param cleanup Function to call after the wrapped function
+ * @returns A function that wraps another function with the context
+ */
+export function createContextWrapper(setup: () => any, cleanup: () => void) {
+  return <T extends (...args: any[]) => any>(fn: T) => {
path: src/js/common/contextManager.ts
user: login: ryanking13
body: Oh, right. Good catch. I wrote `createContextWrapper` first, and extracted the common parts into `withContext`, but I forgot to update the `createContextWrapper` to use `withContext` 😂
created_at: 2025-03-02T01:44:21Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,56 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => any,
+  cleanup: () => void,
+  callback: () => T | Promise<T>,
+): T | Promise<T> {
+  setup();
+  try {
+    const result = callback();
+    if (result instanceof Promise) {
+      return result.finally(() => cleanup());
+    } else {
+      cleanup();
+      return result;
+    }
+  } catch (e) {
+    cleanup();
+    throw e;
+  }
path: src/js/common/contextManager.ts
user: login: ryanking13
body: Looks better, Thanks!
created_at: 2025-03-02T02:01:13Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,56 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => any,
+  cleanup: () => void,
+  callback: () => T | Promise<T>,
+): T | Promise<T> {
path: src/js/common/contextManager.ts
user: login: ryanking13
body: Thanks, I tried this but it seems like TS complains about the return type becuase of `result.finally(...)` block.
created_at: 2025-03-02T02:01:38Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,56 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => any,
+  cleanup: () => void,
+  callback: () => T | Promise<T>,
+): T | Promise<T> {
path: src/js/common/contextManager.ts
user: login: hoodmane
body: You can use `\\ @ts-ignore` or a cast. It's important to give the function a correct type signature for call sites, how to fix the implementation to remove errors is less important.
created_at: 2025-03-02T08:19:40Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,56 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => any,
+  cleanup: () => void,
+  callback: () => T | Promise<T>,
+): T | Promise<T> {
path: src/js/common/contextManager.ts
user: login: ryanking13
body: Sounds reasonable to me. I casted the type in the finally block.
created_at: 2025-03-02T10:46:37Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,50 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => void,
+  cleanup: () => void,
+  callback: () => T,
+): T {
+  setup();
+  let result: T;
+  try {
+    result = callback();
+  } catch (e) {
+    cleanup();
+    throw e;
+  }
+  if (result instanceof Promise) {
+    return result.finally(() => cleanup()) as T;
+  }
+
+  cleanup();
+  return result;
+}
+
+/**
+ * Creates a function wrapper that sets up a context before calling the function
+ * and cleans up afterwards. Works with both synchronous and asynchronous functions.
+ *
+ * @param setup Function to call before the wrapped function
+ * @param cleanup Function to call after the wrapped function
+ * @returns A function that wraps another function with the context
+ */
+export function createContextWrapper(
+  setup: () => void,
+  cleanup: () => void,
+): <T extends (...args: any[]) => any>(
+  fn: T,
+) => (...args: Parameters<T>) => ReturnType<T> {
path: src/js/common/contextManager.ts
user: login: hoodmane
body: Is this different from:
```suggestion
) => T {
```

created_at: 2025-03-02T10:49:23Z
author_association: MEMBER
diff_hunk: @@ -182,12 +183,29 @@ export class PackageManager {
     } = {
       checkIntegrity: true,
     },
+  ): Promise<Array<PackageData>> {
path: src/js/load-package.ts
user: login: hoodmane
body: ```suggestion
  ): Promise<PackageData[]> {
```

created_at: 2025-03-02T11:05:16Z
author_association: MEMBER
diff_hunk: @@ -182,12 +183,29 @@ export class PackageManager {
     } = {
       checkIntegrity: true,
     },
+  ): Promise<Array<PackageData>> {
+    const wrappedLoadPackage = this.setCallbacks(
+      options.messageCallback,
+      options.errorCallback,
+    )(this.loadPackageInner);
+
+    return wrappedLoadPackage.call(this, names, options);
+  }
+
+  public async loadPackageInner(
+    names: string | PyProxy | Array<string>,
path: src/js/load-package.ts
user: login: hoodmane
body: ```suggestion
    names: string | PyProxy | string[],
```

created_at: 2025-03-02T11:05:32Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,50 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => void,
+  cleanup: () => void,
+  callback: () => T,
+): T {
+  setup();
+  let result: T;
+  try {
+    result = callback();
+  } catch (e) {
+    cleanup();
+    throw e;
+  }
+  if (result instanceof Promise) {
+    return result.finally(() => cleanup()) as T;
+  }
+
+  cleanup();
+  return result;
+}
+
+/**
+ * Creates a function wrapper that sets up a context before calling the function
+ * and cleans up afterwards. Works with both synchronous and asynchronous functions.
+ *
+ * @param setup Function to call before the wrapped function
+ * @param cleanup Function to call after the wrapped function
+ * @returns A function that wraps another function with the context
+ */
+export function createContextWrapper(
+  setup: () => void,
+  cleanup: () => void,
+): <T extends (...args: any[]) => any>(
+  fn: T,
+) => (...args: Parameters<T>) => ReturnType<T> {
+  return <T extends (...args: any[]) => any>(fn: T) => {
path: src/js/common/contextManager.ts
user: login: hoodmane
body: I don't think we need to repeat the type since it should be inferred as the return value?
```suggestion
  return (fn) => {
```

created_at: 2025-03-02T11:16:57Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,50 @@
+/**
+ * Python's ContextManager-like higher-order function.
+ * Works with both synchronous and asynchronous callbacks.
+ *
+ * @param setup A function that will be called to setup the context
+ * @param cleanup A function that will be called to cleanup the context
+ * @param callback A function that will be called with the context
+ * @returns The result of the callback function (or a Promise if callback is async)
+ */
+export function withContext<T>(
+  setup: () => void,
+  cleanup: () => void,
+  callback: () => T,
+): T {
+  setup();
+  let result: T;
+  try {
+    result = callback();
+  } catch (e) {
+    cleanup();
+    throw e;
+  }
+  if (result instanceof Promise) {
+    return result.finally(() => cleanup()) as T;
+  }
+
+  cleanup();
+  return result;
+}
+
+/**
+ * Creates a function wrapper that sets up a context before calling the function
+ * and cleans up afterwards. Works with both synchronous and asynchronous functions.
+ *
+ * @param setup Function to call before the wrapped function
+ * @param cleanup Function to call after the wrapped function
+ * @returns A function that wraps another function with the context
+ */
+export function createContextWrapper(
+  setup: () => void,
+  cleanup: () => void,
+): <T extends (...args: any[]) => any>(
+  fn: T,
+) => (...args: Parameters<T>) => ReturnType<T> {
path: src/js/common/contextManager.ts
user: login: ryanking13
body: T refers to the original function, while `(...args: Parameters<T>) => ReturnType<T>` refers to (any) new function that has the same signature as T. So I think they are different.
created_at: 2025-03-02T11:38:20Z
author_association: MEMBER

PR comments:
user: login: ryanking13
created_at: 2025-03-02T11:33:33Z
body: > Let's maybe go though some more rounds of nits on the typescript types but generally looks good to me.

Yeah, thanks for the comments on the types. It is always difficult for me to deal with complex typescript types. I think I have simplified the types enough now. (Claude was way better than me at dealing with types 😂)
author_association: MEMBER
user: login: hoodmane
created_at: 2025-03-03T07:39:54Z
body: Looks great now, thanks!
author_association: MEMBER
user: login: ryanking13
created_at: 2025-03-04T06:14:41Z
body: Okay, I'll go ahead and merge. Thanks for the reviews! There is a failure in `cmdline_runner`, but it is happening in the main branch too. I gues it is related to the recent pytest 8.3.5 release, but we'll need to take a look.
author_association: MEMBER