Correctly handle custom extensions during module resolution (#1400)

Perryvw · web-flow · commit 589baeb7fcd6 · 2023-02-12T21:38:33.000+01:00
* Correctly handle custom extensions during module resolution

* Fix tests
diff --git a/src/transpilation/output-collector.ts b/src/transpilation/output-collector.ts
@@ -14,7 +14,7 @@ export interface TranspiledFile {
     jsSourceMap?: string;
 }
 
-export function createEmitOutputCollector() {
+export function createEmitOutputCollector(luaExtension = ".lua") {
     const files: TranspiledFile[] = [];
     const writeFile: ts.WriteFileCallback = (fileName, data, _bom, _onError, sourceFiles = []) => {
         let file = files.find(f => intersection(f.sourceFiles, sourceFiles).length > 0);
@@ -25,9 +25,9 @@ export function createEmitOutputCollector() {
             file.sourceFiles = union(file.sourceFiles, sourceFiles);
         }
 
-        if (fileName.endsWith(".lua")) {
+        if (fileName.endsWith(luaExtension)) {
             file.lua = data;
-        } else if (fileName.endsWith(".lua.map")) {
+        } else if (fileName.endsWith(`${luaExtension}.map`)) {
             file.luaSourceMap = data;
         } else if (fileName.endsWith(".js")) {
             file.js = data;
diff --git a/src/transpilation/resolve.ts b/src/transpilation/resolve.ts
@@ -49,8 +49,8 @@ class ResolutionContext {
                 // Remove @NoResolution prefix if not building in library mode
                 if (!isBuildModeLibrary(this.program)) {
                     const path = required.requirePath.replace("@NoResolution:", "");
-                    replaceRequireInCode(file, required, path);
-                    replaceRequireInSourceMap(file, required, path);
+                    replaceRequireInCode(file, required, path, this.options.extension);
+                    replaceRequireInSourceMap(file, required, path, this.options.extension);
                 }
 
                 // Skip
@@ -88,8 +88,8 @@ class ResolutionContext {
         // Figure out resolved require path and dependency output path
         if (shouldRewriteRequires(dependencyPath, this.program)) {
             const resolvedRequire = getEmitPathRelativeToOutDir(dependencyPath, this.program);
-            replaceRequireInCode(file, required, resolvedRequire);
-            replaceRequireInSourceMap(file, required, resolvedRequire);
+            replaceRequireInCode(file, required, resolvedRequire, this.options.extension);
+            replaceRequireInSourceMap(file, required, resolvedRequire, this.options.extension);
         }
     }
 
@@ -117,8 +117,8 @@ class ResolutionContext {
 
     private couldNotResolveImport(required: LuaRequire, file: ProcessedFile): void {
         const fallbackRequire = fallbackResolve(required, getSourceDir(this.program), path.dirname(file.fileName));
-        replaceRequireInCode(file, required, fallbackRequire);
-        replaceRequireInSourceMap(file, required, fallbackRequire);
+        replaceRequireInCode(file, required, fallbackRequire, this.options.extension);
+        replaceRequireInSourceMap(file, required, fallbackRequire, this.options.extension);
 
         this.diagnostics.push(
             couldNotResolveRequire(required.requirePath, path.relative(getProjectRoot(this.program), file.fileName))
@@ -310,16 +310,26 @@ function isBuildModeLibrary(program: ts.Program) {
     return program.getCompilerOptions().buildMode === BuildMode.Library;
 }
 
-function replaceRequireInCode(file: ProcessedFile, originalRequire: LuaRequire, newRequire: string): void {
-    const requirePath = formatPathToLuaPath(newRequire.replace(".lua", ""));
+function replaceRequireInCode(
+    file: ProcessedFile,
+    originalRequire: LuaRequire,
+    newRequire: string,
+    extension: string | undefined
+): void {
+    const requirePath = requirePathForFile(newRequire, extension);
     file.code = file.code =
         file.code.substring(0, originalRequire.from) +
         `require("${requirePath}")` +
         file.code.substring(originalRequire.to + 1);
 }
 
-function replaceRequireInSourceMap(file: ProcessedFile, originalRequire: LuaRequire, newRequire: string): void {
-    const requirePath = formatPathToLuaPath(newRequire.replace(".lua", ""));
+function replaceRequireInSourceMap(
+    file: ProcessedFile,
+    originalRequire: LuaRequire,
+    newRequire: string,
+    extension?: string | undefined
+): void {
+    const requirePath = requirePathForFile(newRequire, extension);
     if (file.sourceMapNode) {
         replaceInSourceMap(
             file.sourceMapNode,
@@ -330,6 +340,14 @@ function replaceRequireInSourceMap(file: ProcessedFile, originalRequire: LuaRequ
     }
 }
 
+function requirePathForFile(filePath: string, extension = ".lua"): string {
+    if (filePath.endsWith(extension)) {
+        return formatPathToLuaPath(filePath.substring(0, filePath.length - extension.length));
+    } else {
+        return formatPathToLuaPath(filePath);
+    }
+}
+
 function replaceInSourceMap(node: SourceNode, parent: SourceNode, require: string, resolvedRequire: string): boolean {
     if ((!node.children || node.children.length === 0) && node.toString() === require) {
         parent.children = [new SourceNode(node.line, node.column, node.source, [resolvedRequire])];
diff --git a/test/transpile/module-resolution.spec.ts b/test/transpile/module-resolution.spec.ts
@@ -174,6 +174,14 @@ describe("module resolution with sourceDir", () => {
             .setOptions({ luaBundle: "bundle.lua", luaBundleEntry: mainFile })
             .expectToEqual(expectedResult);
     });
+
+    // https://github.com/TypeScriptToLua/TypeScriptToLua/issues/1394
+    test("can resolve files with non-standard extension (#1394)", () => {
+        util.testProject(path.join(projectPath, "tsconfig.json"))
+            .setMainFileName(path.join(projectPath, "src", "main.ts"))
+            .setOptions({ outDir: "tstl-out", extension: ".script" })
+            .expectToEqual(expectedResult);
+    });
 });
 
 describe("module resolution project with lua sources", () => {
diff --git a/test/util.ts b/test/util.ts
@@ -236,7 +236,7 @@ export abstract class TestBuilder {
     public getLuaResult(): tstl.TranspileVirtualProjectResult {
         const program = this.getProgram();
         const preEmitDiagnostics = ts.getPreEmitDiagnostics(program);
-        const collector = createEmitOutputCollector();
+        const collector = createEmitOutputCollector(this.options.extension);
         const { diagnostics: transpileDiagnostics } = new tstl.Transpiler({ emitHost: this.getEmitHost() }).emit({
             program,
             customTransformers: this.customTransformers,
@@ -277,7 +277,7 @@ export abstract class TestBuilder {
         const program = this.getProgram();
         program.getCompilerOptions().module = ts.ModuleKind.CommonJS;
 
-        const collector = createEmitOutputCollector();
+        const collector = createEmitOutputCollector(this.options.extension);
         const { diagnostics } = program.emit(undefined, collector.writeFile);
         return { transpiledFiles: collector.files, diagnostics: [...diagnostics] };
     }
@@ -468,7 +468,10 @@ end)());`;
     }
 
     private injectLuaFile(state: LuaState, lua: Lua, lauxlib: LauxLib, fileName: string, fileContent: string) {
-        const modName = formatPathToLuaPath(fileName.replace(".lua", ""));
+        const extension = this.options.extension ?? ".lua";
+        const modName = fileName.endsWith(extension)
+            ? formatPathToLuaPath(fileName.substring(0, fileName.length - extension.length))
+            : fileName;
         if (this.options.luaTarget === tstl.LuaTarget.Lua50) {
             // Adding source Lua to the _LOADED cache will allow require to find it
             lua.lua_getglobal(state, "_LOADED");
@@ -610,7 +613,7 @@ class ProjectTestBuilder extends ModuleTestBuilder {
     @memoize
     public getLuaResult(): tstl.TranspileVirtualProjectResult {
         // Override getLuaResult to use transpileProject with tsconfig.json instead
-        const collector = createEmitOutputCollector();
+        const collector = createEmitOutputCollector(this.options.extension);
         const { diagnostics } = transpileProject(this.tsConfig, this.options, collector.writeFile);
 
         return { diagnostics: [...diagnostics], transpiledFiles: collector.files };
