From 1ad76e0d6f5851cda82978a8081ceb85e9c92ca5 Mon Sep 17 00:00:00 2001
From: calebboyd <caleb.boyd@gmail.com>
Date: Thu, 3 Jun 2021 12:37:07 -0500
Subject: [PATCH] feat(register): add try-ts-ext register entrypoint

---
 package.json           |  1 +
 register/try-ts-ext.js |  3 +++
 src/index.ts           | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+)
 create mode 100644 register/try-ts-ext.js

diff --git a/package.json b/package.json
index 76dbbe56f..7cb0a54b3 100644
--- a/package.json
+++ b/package.json
@@ -19,6 +19,7 @@
     "./register/files": "./register/files.js",
     "./register/transpile-only": "./register/transpile-only.js",
     "./register/type-check": "./register/type-check.js",
+    "./register/try-ts-ext": "./register/try-ts-ext.js",
     "./esm": "./esm.mjs",
     "./esm.mjs": "./esm.mjs",
     "./esm/transpile-only": "./esm/transpile-only.mjs",
diff --git a/register/try-ts-ext.js b/register/try-ts-ext.js
new file mode 100644
index 000000000..e2d59618b
--- /dev/null
+++ b/register/try-ts-ext.js
@@ -0,0 +1,3 @@
+require('../dist').register({
+  tryTsExt: true,
+});
diff --git a/src/index.ts b/src/index.ts
index 9aea226d4..b6da2d8d2 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -84,6 +84,7 @@ export interface ProcessEnv {
   TS_NODE_SKIP_PROJECT?: string;
   TS_NODE_SKIP_IGNORE?: string;
   TS_NODE_PREFER_TS_EXTS?: string;
+  TS_NODE_TRY_TS_EXT?: string;
   TS_NODE_IGNORE_DIAGNOSTICS?: string;
   TS_NODE_TRANSPILE_ONLY?: string;
   TS_NODE_TYPE_CHECK?: string;
@@ -280,6 +281,13 @@ export interface RegisterOptions extends CreateOptions {
    * @default false
    */
   preferTsExts?: boolean;
+
+  /**
+   * Attempt to resolve the typescript file when the js file cannot be found.
+   *
+   * @default false
+   */
+  tryTsExt?: boolean
 }
 
 /**
@@ -327,6 +335,7 @@ export const DEFAULTS: RegisterOptions = {
   skipProject: yn(env.TS_NODE_SKIP_PROJECT),
   skipIgnore: yn(env.TS_NODE_SKIP_IGNORE),
   preferTsExts: yn(env.TS_NODE_PREFER_TS_EXTS),
+  tryTsExt: yn(env.TS_NODE_TRY_TS_EXT),
   ignoreDiagnostics: split(env.TS_NODE_IGNORE_DIAGNOSTICS),
   transpileOnly: yn(env.TS_NODE_TRANSPILE_ONLY),
   typeCheck: yn(env.TS_NODE_TYPE_CHECK),
@@ -402,6 +411,27 @@ export function getExtensions(config: _ts.ParsedCommandLine) {
   return { tsExtensions, jsExtensions };
 }
 
+function patchResolveFileName() {
+  const originalResolveFilename = (Module as any)._resolveFilename;
+  const jsExt = /\.[mc]?js(x)?/;
+
+  (Module as any)._resolveFilename = function (...args: any[]) {
+    const [request, _, isMain ] = args
+    if (isMain) {
+      return originalResolveFilename.apply(this, args);
+    }
+    try {
+      return originalResolveFilename.apply(this, args);
+    } catch (e) {
+      if (e.code === 'MODULE_NOT_FOUND' && jsExt.test(request)) {
+        args[0] = request.replace(jsExt, '.ts$1');
+        return originalResolveFilename.apply(this, args);
+      }
+      throw e;
+    }
+  };
+}
+
 /**
  * Register TypeScript compiler instance onto node.js
  */
@@ -422,6 +452,10 @@ export function register(opts: RegisterOptions = {}): Service {
     originalJsHandler
   );
 
+  if (service.options.tryTsExt) {
+    patchResolveFileName();
+  }
+
   // Require specified modules before start-up.
   (Module as any)._preloadModules(service.options.require);