Refactor Hooks

packages/engine/src/CellEvaluator.ts

@@ -1,5 +1,6 @@
 import { TypeCellContext } from "./context";
 import { ModuleExecution, runModule } from "./executor";
+import { HookExecution } from "./HookExecution";
 import { createExecutionScope, getModulesFromTypeCellCode } from "./modules";
 import { isReactView } from "./reactView";
 
@@ -49,7 +50,11 @@ export function createCellEvaluator(
     onOutputChanged(error);
   }
 
-  const executionScope = createExecutionScope(typecellContext);
+  const hookExecution = new HookExecution();
+  const executionScope = createExecutionScope(
+    typecellContext,
+    hookExecution.context
+  );
   let moduleExecution: ModuleExecution | undefined;
 
   async function evaluate(compiledCode: string) {
@@ -69,6 +74,7 @@ export function createCellEvaluator(
         modules[0],
         typecellContext,
         resolveImport,
+        hookExecution,
         beforeExecuting,
         onExecuted,
         onError,

packages/engine/src/HookExecution.ts

@@ -0,0 +1,97 @@
+const glob = typeof window === "undefined" ? global : window;
+
+const overrideFunctions = [
+  "setTimeout",
+  "setInterval",
+  "console",
+  "EventTarget.prototype.addEventListener",
+] as const;
+
+const originalReferences: HookContext = {
+  setTimeout: glob.setTimeout,
+  setInterval: glob.setInterval,
+  console: glob.console,
+  "EventTarget.prototype.addEventListener":
+    glob.EventTarget.prototype.addEventListener,
+};
+
+export type HookContext = { [K in typeof overrideFunctions[number]]: any };
+
+function setProperty(base: Object, path: string, value: any) {
+  const layers = path.split(".");
+  if (layers.length > 1) {
+    const toOverride = layers.pop()!;
+    // Returns second last path member
+    const layer = layers.reduce((o, i) => o[i], base as any);
+    layer[toOverride] = value;
+  } else {
+    (base as any)[path] = value;
+  }
+}
+
+export class HookExecution {
+  public disposers: Array<() => void> = [];
+  public context: HookContext = {
+    setTimeout: this.createHookedFunction(setTimeout, this.disposers, (ret) => {
+      clearTimeout(ret);
+    }),
+    setInterval: this.createHookedFunction(
+      setInterval,
+      this.disposers,
+      (ret) => {
+        clearInterval(ret);
+      }
+    ),
+    console: {
+      ...originalReferences.console,
+      log: (...args: any) => {
+        // TODO: broadcast output to console view
+        originalReferences.console.log(...args);
+      },
+    },
+    ["EventTarget.prototype.addEventListener"]: undefined,
+  };
+
+  constructor() {
+    if (typeof EventTarget !== "undefined") {
+      this.context["EventTarget.prototype.addEventListener"] =
+        this.createHookedFunction(
+          EventTarget.prototype.addEventListener as any,
+          this.disposers,
+          function (this: any, _ret, args) {
+            this.removeEventListener(args[0], args[1]);
+          }
+        );
+    }
+  }
+
+  public attachToWindow() {
+    overrideFunctions.forEach((path) =>
+      setProperty(glob, path, this.context[path])
+    );
+  }
+
+  public detachFromWindow() {
+    overrideFunctions.forEach((path) =>
+      setProperty(glob, path, originalReferences[path])
+    );
+  }
+
+  public dispose() {
+    this.disposers.forEach((d) => d());
+  }
+
+  private createHookedFunction<T, Y>(
+    original: (...args: T[]) => Y,
+    disposes: Array<() => void>,
+    disposer: (ret: Y, args: T[]) => void
+  ) {
+    return function newFunction(this: any): Y {
+      const callerArguments = arguments;
+      const ret = original.apply(this, callerArguments as any); // TODO: fix any?
+      const ctx = this;
+      disposes.push(() => disposer.call(ctx, ret, callerArguments as any));
+      return ret;
+    };
+  }
+}

packages/engine/src/executor.ts

@@ -1,6 +1,6 @@
 import { autorun, runInAction } from "mobx";
 import { TypeCellContext } from "./context";
-import { installHooks } from "./hookDisposables";
+import { HookExecution } from "./HookExecution";
 import { Module } from "./modules";
 import { isStored } from "./storage/stored";
 import { isView } from "./view";
@@ -61,6 +61,7 @@ export async function runModule(
   mod: Module,
   context: TypeCellContext<any>,
   resolveImport: (module: string) => any,
+  hookExecution: HookExecution,
   beforeExecuting: () => void,
   onExecuted: (exports: any) => void,
   onError: (error: any) => void,
@@ -96,7 +97,7 @@ export async function runModule(
     );
   }
 
-  const execute = async () => {
+  async function execute(this: any) {
     try {
       if (wouldLoopOnAutorun) {
         detectedLoop = true;
@@ -114,25 +115,20 @@ export async function runModule(
       disposeEveryRun.length = 0; // clear existing array in this way, because we've passed the reference to resolveDependencyArray and want to keep it intact
 
       beforeExecuting();
-      const hooks = installHooks();
-      disposeEveryRun.push(hooks.disposeAll);
-      let executionPromise: Promise<any>;
+
+      disposeEveryRun.push(hookExecution.dispose.bind(hookExecution));
+      hookExecution.attachToWindow();
+
       try {
-        executionPromise = mod.factoryFunction.apply(
-          undefined,
-          argsToCallFunctionWith
-        ); // TODO: what happens with disposers if a rerun of this function is slow / delayed?
+        await mod.factoryFunction.apply(undefined, argsToCallFunctionWith);
       } finally {
-        // Hooks are only installed for sync code. Ideally, we'd want to run it for all code, but then we have the chance hooks will affect other parts of the TypeCell (non-user) code
-        // (we ran into this that notebooks wouldn't be saved (_.debounce), and also that setTimeout of Monaco blink cursor would be hooked)
-        hooks.unHookAll();
+        hookExecution.detachFromWindow();
+
         if (previousVariableDisposer) {
           previousVariableDisposer(exports);
         }
       }
 
-      await executionPromise;
-
       // Running the assignments to `context` in action should be a performance improvement to prevent triggering observers one-by-one
       wouldLoopOnAutorun = true;
       runInAction(() => {
@@ -211,7 +207,7 @@ export async function runModule(
       //reject(e);
       resolve();
     }
-  };
+  }
 
   const autorunDisposer = autorun(() => execute());
 

packages/engine/src/modules.ts

@@ -1,5 +1,7 @@
 import { TypeCellContext } from "./context";
 import { observable, untracked, computed, autorun } from "mobx";
+import { HookContext } from "./HookExecution";
+import { set as setProperty } from "lodash";
 // import { stored } from "./storage/stored";
 // import { view } from "./view";
 
@@ -71,7 +73,10 @@ function createDefine(modules: Module[]) {
   };
 }
 
-export function createExecutionScope(context: TypeCellContext<any>) {
+export function createExecutionScope(
+  context: TypeCellContext<any>,
+  hookContext: HookContext
+) {
   const scope = {
     autorun,
     $: context.context,
@@ -83,6 +88,11 @@ export function createExecutionScope(context: TypeCellContext<any>) {
     // view,
     observable,
   };
+
+  Object.keys(hookContext).forEach((path) =>
+    setProperty(scope, path, hookContext[path as keyof HookContext])
+  );
+
   return scope;
 }