From bb2e9f1353697974fff1401146655c95ff32ef3c Mon Sep 17 00:00:00 2001
From: isaacs <i@izs.me>
Date: Tue, 5 Sep 2023 17:32:52 -0700
Subject: [PATCH] add support for '--import ts-node/import'

This also adds a type for the loader hooks API v3, as globalPreload is
scheduled for removal in node v21, at which point '--loader ts-node/esm'
will no longer work, and '--import ts-node/import' will be the way
forward.
---
 README.md  |  2 +-
 esm.mjs    |  2 +-
 src/esm.ts | 39 +++++++++++++++++++++++++++++----------
 3 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 40a40947f..cb2fb9981 100644
--- a/README.md
+++ b/README.md
@@ -230,7 +230,7 @@ To test your version of `env` for compatibility with `-S`:
 
 ## node flags and other tools
 
-You can register ts-node without using our CLI: `node -r ts-node/register` and `node --loader ts-node/esm`
+You can register ts-node without using our CLI: `node -r ts-node/register`, `node --loader ts-node/esm`, or `node --import ts-node/import` in node 20.6 and above.
 
 In many cases, setting [`NODE_OPTIONS`](https://nodejs.org/api/cli.html#cli_node_options_options) will enable `ts-node` within other node tools, child processes, and worker threads.  This can be combined with other node flags.
 
diff --git a/esm.mjs b/esm.mjs
index 4a4b2e323..07ed7008a 100644
--- a/esm.mjs
+++ b/esm.mjs
@@ -4,4 +4,4 @@ const require = createRequire(fileURLToPath(import.meta.url));
 
 /** @type {import('./dist/esm')} */
 const esm = require('./dist/esm');
-export const { resolve, load, getFormat, transformSource, globalPreload } = esm.registerAndCreateEsmHooks();
+export const { initialize, resolve, load, getFormat, transformSource, globalPreload } = esm.registerAndCreateEsmHooks();
diff --git a/src/esm.ts b/src/esm.ts
index c4566e0df..8de6fb075 100644
--- a/src/esm.ts
+++ b/src/esm.ts
@@ -79,6 +79,17 @@ export namespace NodeLoaderHooksAPI2 {
   export type GlobalPreloadHook = (context?: { port: MessagePort }) => string;
 }
 
+export interface NodeLoaderHooksAPI3 {
+  resolve: NodeLoaderHooksAPI2.ResolveHook;
+  load: NodeLoaderHooksAPI2.LoadHook;
+  initialize?: NodeLoaderHooksAPI3.InitializeHook;
+}
+export namespace NodeLoaderHooksAPI3 {
+  // technically this can be anything that can be passed through a postMessage channel,
+  // but defined here based on how ts-node uses it.
+  export type InitializeHook = (data: any) => void | Promise<void>;
+}
+
 export type NodeLoaderHooksFormat = 'builtin' | 'commonjs' | 'dynamic' | 'json' | 'module' | 'wasm';
 
 export type NodeImportConditions = unknown;
@@ -87,17 +98,24 @@ export interface NodeImportAssertions {
 }
 
 // The hooks API changed in node version X so we need to check for backwards compatibility.
-const newHooksAPI = versionGteLt(process.versions.node, '16.12.0');
+const hooksAPIVersion = versionGteLt(process.versions.node, '21.0.0')
+  ? 3
+  : versionGteLt(process.versions.node, '16.12.0')
+  ? 2
+  : 1;
 
 /** @internal */
 export function filterHooksByAPIVersion(
-  hooks: NodeLoaderHooksAPI1 & NodeLoaderHooksAPI2
-): NodeLoaderHooksAPI1 | NodeLoaderHooksAPI2 {
-  const { getFormat, load, resolve, transformSource, globalPreload } = hooks;
+  hooks: NodeLoaderHooksAPI1 & NodeLoaderHooksAPI2 & NodeLoaderHooksAPI3
+): NodeLoaderHooksAPI1 | NodeLoaderHooksAPI2 | NodeLoaderHooksAPI3 {
+  const { getFormat, load, resolve, transformSource, globalPreload, initialize } = hooks;
   // Explicit return type to avoid TS's non-ideal inferred type
-  const hooksAPI: NodeLoaderHooksAPI1 | NodeLoaderHooksAPI2 = newHooksAPI
-    ? { resolve, load, globalPreload, getFormat: undefined, transformSource: undefined }
-    : { resolve, getFormat, transformSource, load: undefined };
+  const hooksAPI: NodeLoaderHooksAPI1 | NodeLoaderHooksAPI2 | NodeLoaderHooksAPI3 =
+    hooksAPIVersion === 3
+      ? { resolve, load, initialize, globalPreload: undefined, transformSource: undefined, getFormat: undefined }
+      : hooksAPIVersion === 2
+      ? { resolve, load, globalPreload, initialize: undefined, getFormat: undefined, transformSource: undefined }
+      : { resolve, getFormat, transformSource, initialize: undefined, globalPreload: undefined, load: undefined };
   return hooksAPI;
 }
 
@@ -122,6 +140,7 @@ export function createEsmHooks(tsNodeService: Service) {
     getFormat,
     transformSource,
     globalPreload: useLoaderThread ? globalPreload : undefined,
+    initialize: undefined,
   });
 
   function globalPreload({ port }: { port?: MessagePort } = {}) {
@@ -129,9 +148,9 @@ export function createEsmHooks(tsNodeService: Service) {
     // so this signal lets us infer it based on the state of the main
     // thread, but only relevant if options.pretty is unset.
     let stderrTTYSignal: string;
-    if (tsNodeService.options.pretty === undefined) {
-      port?.on('message', (data) => {
-        if (data?.stderrIsTTY) {
+    if (port && tsNodeService.options.pretty === undefined) {
+      port.on('message', (data: { stderrIsTTY?: boolean }) => {
+        if (data.stderrIsTTY) {
           tsNodeService.setPrettyErrors(true);
         }
       });