1 files changed, 30 insertions, 20 deletions

diff --git a/src/client/util/Scripting.ts b/src/client/util/Scripting.ts
index 46dc320b0..1d0916ac0 100644
--- a/src/client/util/Scripting.ts
+++ b/src/client/util/Scripting.ts
@@ -52,10 +52,10 @@ export namespace Scripting {
         } else {
             throw new Error("Must either register an object with a name, or give a name and an object");
         }
-        if (scriptingGlobals.hasOwnProperty(n)) {
+        if (_scriptingGlobals.hasOwnProperty(n)) {
             throw new Error(`Global with name ${n} is already registered, choose another name`);
         }
-        scriptingGlobals[n] = obj;
+        _scriptingGlobals[n] = obj;
     }
 
     export function makeMutableGlobalsCopy(globals?: { [name: string]: any }) {
@@ -188,6 +188,10 @@ class ScriptingCompilerHost {
 
 export type Traverser = (node: ts.Node, indentation: string) => boolean | void;
 export type TraverserParam = Traverser | { onEnter: Traverser, onLeave: Traverser };
+export type Transformer = {
+    transformer: ts.TransformerFactory<ts.SourceFile>,
+    getVars?: () => { capturedVariables: { [name: string]: Field } }
+};
 export interface ScriptOptions {
     requiredType?: string;
     addReturn?: boolean;
@@ -196,7 +200,7 @@ export interface ScriptOptions {
     typecheck?: boolean;
     editable?: boolean;
     traverser?: TraverserParam;
-    transformer?: ts.TransformerFactory<ts.SourceFile>;
+    transformer?: Transformer;
     globals?: { [name: string]: any };
 }
 
@@ -213,6 +217,27 @@ export function CompileScript(script: string, options: ScriptOptions = {}): Comp
         Scripting.setScriptingGlobals(options.globals);
     }
     let host = new ScriptingCompilerHost;
+    if (options.traverser) {
+        const sourceFile = ts.createSourceFile('script.ts', script, ts.ScriptTarget.ES2015, true);
+        const onEnter = typeof options.traverser === "object" ? options.traverser.onEnter : options.traverser;
+        const onLeave = typeof options.traverser === "object" ? options.traverser.onLeave : undefined;
+        forEachNode(sourceFile, onEnter, onLeave);
+    }
+    if (options.transformer) {
+        const sourceFile = ts.createSourceFile('script.ts', script, ts.ScriptTarget.ES2015, true);
+        const result = ts.transform(sourceFile, [options.transformer.transformer]);
+        if (options.transformer.getVars) {
+            const newCaptures = options.transformer.getVars();
+            // tslint:disable-next-line: prefer-object-spread
+            options.capturedVariables = Object.assign(capturedVariables, newCaptures.capturedVariables) as any;
+        }
+        const transformed = result.transformed;
+        const printer = ts.createPrinter({
+            newLine: ts.NewLineKind.LineFeed
+        });
+        script = printer.printFile(transformed[0]);
+        result.dispose();
+    }
     let paramNames: string[] = [];
     if ("this" in params || "this" in capturedVariables) {
         paramNames.push("this");
@@ -227,26 +252,11 @@ export function CompileScript(script: string, options: ScriptOptions = {}): Comp
     });
     for (const key in capturedVariables) {
         if (key === "this") continue;
+        const val = capturedVariables[key];
         paramNames.push(key);
-        paramList.push(`${key}: ${capturedVariables[key].constructor.name}`);
+        paramList.push(`${key}: ${typeof val === "object" ? Object.getPrototypeOf(val).constructor.name : typeof val}`);
     }
     let paramString = paramList.join(", ");
-    if (options.traverser) {
-        const sourceFile = ts.createSourceFile('script.ts', script, ts.ScriptTarget.ES2015, true);
-        const onEnter = typeof options.traverser === "object" ? options.traverser.onEnter : options.traverser;
-        const onLeave = typeof options.traverser === "object" ? options.traverser.onLeave : undefined;
-        forEachNode(sourceFile, onEnter, onLeave);
-    }
-    if (options.transformer) {
-        const sourceFile = ts.createSourceFile('script.ts', script, ts.ScriptTarget.ES2015, true);
-        const result = ts.transform(sourceFile, [options.transformer]);
-        const transformed = result.transformed;
-        const printer = ts.createPrinter({
-            newLine: ts.NewLineKind.LineFeed
-        });
-        script = printer.printFile(transformed[0]);
-        result.dispose();
-    }
     let funcScript = `(function(${paramString})${requiredType ? `: ${requiredType}` : ''} {
         ${addReturn ? `return ${script};` : script}
     })`;