Make printer slightly easier to use, add better printer example as test (#1173)

Perryvw · web-flow · commit 25dd8c83d309 · 2021-11-23T13:36:13.000+01:00
diff --git a/src/LuaPrinter.ts b/src/LuaPrinter.ts
@@ -158,9 +158,9 @@ export class LuaPrinter {
     private static rightAssociativeOperators = new Set([lua.SyntaxKind.ConcatOperator, lua.SyntaxKind.PowerOperator]);
 
     private currentIndent = "";
-    private luaFile: string;
-    private relativeSourcePath: string;
-    private options: CompilerOptions;
+    protected luaFile: string;
+    protected relativeSourcePath: string;
+    protected options: CompilerOptions;
 
     public static readonly sourceMapTracebackPlaceholder = "{#SourceMapTraceback}";
 
@@ -226,7 +226,7 @@ export class LuaPrinter {
         return `__TS__SourceMapTraceBack(debug.getinfo(1).short_src, ${mapString});`;
     }
 
-    private printFile(file: lua.File): SourceNode {
+    protected printFile(file: lua.File): SourceNode {
         let header = file.trivia;
 
         if (!this.options.noHeader) {
diff --git a/test/transpile/plugins/plugins.spec.ts b/test/transpile/plugins/plugins.spec.ts
@@ -4,7 +4,30 @@ import * as util from "../../util";
 test("printer", () => {
     util.testModule``
         .setOptions({ luaPlugins: [{ name: path.join(__dirname, "printer.ts") }] })
-        .tap(builder => expect(builder.getMainLuaCodeChunk()).toMatch("Plugin"));
+        .tap(builder => expect(builder.getMainLuaCodeChunk()).toMatch("-- Custom printer plugin:"));
+});
+
+test("printer in bundle", () => {
+    const { transpiledFiles } = util.testBundle`
+        import "./otherfile";
+
+        const foo = "foo";
+    `
+        .addExtraFile(
+            "otherfile.ts",
+            `
+                const bar = "bar";
+            `
+        )
+        .setOptions({ luaPlugins: [{ name: path.join(__dirname, "printer.ts") }] })
+        .expectToHaveNoDiagnostics()
+        .getLuaResult();
+
+    expect(transpiledFiles).toHaveLength(1);
+    const lua = transpiledFiles[0].lua!;
+
+    expect(lua).toContain("-- Custom printer plugin: main.lua");
+    expect(lua).toContain("-- Custom printer plugin: otherfile.lua");
 });
 
 test("visitor", () => {
diff --git a/test/transpile/plugins/printer.ts b/test/transpile/plugins/printer.ts
@@ -1,11 +1,23 @@
+import { SourceNode } from "source-map";
 import * as tstl from "../../../src";
 
+class CustomPrinter extends tstl.LuaPrinter {
+    /* Override printFile */
+    protected printFile(file: tstl.File): SourceNode {
+        const originalResult = super.printFile(file);
+        // Add header comment at the top of the file
+        return this.createSourceNode(file, [`-- Custom printer plugin: ${this.luaFile}\n`, originalResult]);
+    }
+
+    /* Override printBoolean */
+    public printBooleanLiteral(expression: tstl.BooleanLiteral): SourceNode {
+        // Print any boolean as 'true'
+        return this.createSourceNode(expression, "true");
+    }
+}
+
 const plugin: tstl.Plugin = {
-    printer(program, emitHost, fileName, ...args) {
-        const result = new tstl.LuaPrinter(emitHost, program, fileName).print(...args);
-        result.code = `-- Plugin\n${result.code}`;
-        return result;
-    },
+    printer: (program, emitHost, fileName, file) => new CustomPrinter(emitHost, program, fileName).print(file),
 };
 
 // eslint-disable-next-line import/no-default-export
