From 248fdc24a61aca946902b20b9d8187101a9b7bfa Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Fri, 18 Oct 2024 12:26:57 +0200
Subject: [PATCH] feat(typescript-bun): support relative imports without the
 .ts extension

---
 backend/windmill-worker/loader.bun.js | 69 ++++++++++++++++++++-------
 frontend/src/lib/ata/index.ts         | 11 +++--
 2 files changed, 61 insertions(+), 19 deletions(-)

diff --git a/backend/windmill-worker/loader.bun.js b/backend/windmill-worker/loader.bun.js
index b6efb68774dac..ce75cbfc15b40 100644
--- a/backend/windmill-worker/loader.bun.js
+++ b/backend/windmill-worker/loader.bun.js
@@ -8,29 +8,65 @@ const p = {
       "localhost",
       "127.0.0.1"
     );
+
     const w_id = "W_ID";
     const current_path = "CURRENT_PATH";
+    const token = "TOKEN";
+
+    const cdir = resolve("./");
+    const cdirNoPrivate = cdir.replace(/^\/private/, ""); // for macos
+    const filterResolve = new RegExp(
+      `^(?!\\.\/main\\.ts)(?!${cdir}\/main\\.ts)(?!(?:/private)?${cdirNoPrivate}\/wrapper\\.mjs).*\\.ts$`
+    );
+
+    let cdirNodeModules = `${cdir}/node_modules/`;
+
+    const filterLoad = new RegExp(`^${cdir}\/main\\.ts$`);
+    const transpiler = new Bun.Transpiler({
+      loader: "tsx",
+    });
+
+    function replaceRelativeImports(code) {
+      const imports = transpiler.scanImports(code);
+      for (const imp of imports) {
+        if (imp.kind == "import-statement") {
+          if (imp.path.startsWith(".") && !imp.path.endsWith(".ts")) {
+            code = code.replaceAll(imp.path, imp.path + ".ts");
+          }
+        }
+      }
+      return {
+        contents: code,
+      };
+    }
+
+    build.onLoad({ filter: filterLoad }, async (args) => {
+      const code = readFileSync(args.path, "utf8");
+      return replaceRelativeImports(code);
+    });
 
     build.onLoad({ filter: /.*\.url$/ }, async (args) => {
       const url = readFileSync(args.path, "utf8");
-      const contents = await (
-        await fetch(url, {
-          method: "GET",
-          headers: { Authorization: "Bearer TOKEN" },
-        })
-      ).text();
+      const req = await fetch(url, {
+        method: "GET",
+        headers: {
+          Authorization: "Bearer " + token,
+        },
+      });
+      if (!req.ok) {
+        throw new Error(
+          `Failed to find relative import at ${url}`,
+          req.statusText
+        );
+      }
+      const contents = await req.text();
       return {
-        contents,
+        contents: replaceRelativeImports(contents).contents,
         loader: "tsx",
       };
     });
-    const cdir = resolve("./");
-    const cdirNoPrivate = cdir.replace(/^\/private/, ""); // for macos
-    const filter = new RegExp(
-      `^(?!\\.\/main\\.ts)(?!${cdir}\/main\\.ts)(?!(?:/private)?${cdirNoPrivate}\/wrapper\\.mjs).*\\.ts$`
-    );
-    let cdirNodeModules = `${cdir}/node_modules/`;
-    build.onResolve({ filter }, (args) => {
+
+    build.onResolve({ filter: filterResolve }, (args) => {
       if (args.importer?.startsWith(cdirNodeModules)) {
         return undefined;
       }
@@ -41,9 +77,10 @@ const p = {
 
       const isRelative = !args.path.startsWith("/");
 
+      let endExt = args.path.endsWith(".ts") ? "" : ".ts";
       const url = isRelative
-        ? `${base_internal_url}/api/w/${w_id}/scripts/RAW_GET_ENDPOINT/p/${file_path}/../${args.path}`
-        : `${base_internal_url}/api/w/${w_id}/scripts/RAW_GET_ENDPOINT/p/${args.path}`;
+        ? `${base_internal_url}/api/w/${w_id}/scripts/raw_unpinned/p/${file_path}/../${args.path}${endExt}`
+        : `${base_internal_url}/api/w/${w_id}/scripts/raw_unpinned/p/${args.path}${endExt}`;
       const file = isRelative
         ? resolve("./" + file_path + "/../" + args.path + ".url")
         : resolve("./" + args.path + ".url");
diff --git a/frontend/src/lib/ata/index.ts b/frontend/src/lib/ata/index.ts
index 3af147b45c466..5414df8b186f4 100644
--- a/frontend/src/lib/ata/index.ts
+++ b/frontend/src/lib/ata/index.ts
@@ -120,11 +120,16 @@ export const setupTypeAcquisition = (config: ATABootstrapConfig) => {
 					? f.raw
 					: '/' + config.scriptPath + (f.raw.startsWith('../') ? '/../' : '/.') + f.raw
 				let url = config.root + path
-				// console.log('FOO', config.scriptPath, path, f.raw)
-				console.log('fetching local file', url, f.raw)
+				let localPath = f.raw
+				if (f.raw.startsWith('.') && !f.raw.endsWith('.ts')) {
+					url += '.ts'
+					localPath += '.ts'
+				}
+
+				console.log('fetching local file', url, f.raw, localPath)
 				const res = await fetch(url)
 				if (res.ok) {
-					config.delegate.localFile?.(await res.text(), f.raw)
+					config.delegate.localFile?.(await res.text(), localPath)
 				}
 			})
 		}