From faae24356a4cee3a19fb6faa58b7b269fd11f62e Mon Sep 17 00:00:00 2001
From: Tony Brix <tony@brix.ninja>
Date: Mon, 11 Dec 2023 23:02:22 -0700
Subject: [PATCH] feat: add processAllTokens hook (#3114)

---
 src/Hooks.ts            | 11 ++++++-
 src/Instance.ts         | 16 +++++++---
 src/MarkedOptions.ts    | 21 ++++++++----
 test/types/marked.ts    | 21 ++++++++++++
 test/unit/Hooks.test.js | 71 ++++++++++++++++++++++++++++++++++++++++-
 5 files changed, 126 insertions(+), 14 deletions(-)

diff --git a/src/Hooks.ts b/src/Hooks.ts
index 1a7a1d385f..49a80e8f65 100644
--- a/src/Hooks.ts
+++ b/src/Hooks.ts
@@ -1,5 +1,6 @@
 import { _defaults } from './defaults.ts';
 import type { MarkedOptions } from './MarkedOptions.ts';
+import type { Token, TokensList } from './Tokens.ts';
 
 export class _Hooks {
   options: MarkedOptions;
@@ -10,7 +11,8 @@ export class _Hooks {
 
   static passThroughHooks = new Set([
     'preprocess',
-    'postprocess'
+    'postprocess',
+    'processAllTokens'
   ]);
 
   /**
@@ -26,4 +28,11 @@ export class _Hooks {
   postprocess(html: string) {
     return html;
   }
+
+  /**
+   * Process all tokens before walk tokens
+   */
+  processAllTokens(tokens: Token[] | TokensList) {
+    return tokens;
+  }
 }
diff --git a/src/Instance.ts b/src/Instance.ts
index bd0553ba16..2f8c7b6425 100644
--- a/src/Instance.ts
+++ b/src/Instance.ts
@@ -204,23 +204,25 @@ export class Marked {
           const hooksFunc = pack.hooks[hooksProp] as UnknownFunction;
           const prevHook = hooks[hooksProp] as UnknownFunction;
           if (_Hooks.passThroughHooks.has(prop)) {
-            hooks[hooksProp] = (arg: string | undefined) => {
+            // @ts-expect-error cannot type hook function dynamically
+            hooks[hooksProp] = (arg: unknown) => {
               if (this.defaults.async) {
                 return Promise.resolve(hooksFunc.call(hooks, arg)).then(ret => {
-                  return prevHook.call(hooks, ret) as string;
+                  return prevHook.call(hooks, ret);
                 });
               }
 
               const ret = hooksFunc.call(hooks, arg);
-              return prevHook.call(hooks, ret) as string;
+              return prevHook.call(hooks, ret);
             };
           } else {
+            // @ts-expect-error cannot type hook function dynamically
             hooks[hooksProp] = (...args: unknown[]) => {
               let ret = hooksFunc.apply(hooks, args);
               if (ret === false) {
                 ret = prevHook.apply(hooks, args);
               }
-              return ret as string;
+              return ret;
             };
           }
         }
@@ -292,6 +294,7 @@ export class Marked {
       if (opt.async) {
         return Promise.resolve(opt.hooks ? opt.hooks.preprocess(src) : src)
           .then(src => lexer(src, opt))
+          .then(tokens => opt.hooks ? opt.hooks.processAllTokens(tokens) : tokens)
           .then(tokens => opt.walkTokens ? Promise.all(this.walkTokens(tokens, opt.walkTokens)).then(() => tokens) : tokens)
           .then(tokens => parser(tokens, opt))
           .then(html => opt.hooks ? opt.hooks.postprocess(html) : html)
@@ -302,7 +305,10 @@ export class Marked {
         if (opt.hooks) {
           src = opt.hooks.preprocess(src) as string;
         }
-        const tokens = lexer(src, opt);
+        let tokens = lexer(src, opt);
+        if (opt.hooks) {
+          tokens = opt.hooks.processAllTokens(tokens) as Token[] | TokensList;
+        }
         if (opt.walkTokens) {
           this.walkTokens(tokens, opt.walkTokens);
         }
diff --git a/src/MarkedOptions.ts b/src/MarkedOptions.ts
index e4370501ef..78754f23fe 100644
--- a/src/MarkedOptions.ts
+++ b/src/MarkedOptions.ts
@@ -3,6 +3,7 @@ import type { _Parser } from './Parser.ts';
 import type { _Lexer } from './Lexer.ts';
 import type { _Renderer } from './Renderer.ts';
 import type { _Tokenizer } from './Tokenizer.ts';
+import type { _Hooks } from './Hooks.ts';
 
 export interface TokenizerThis {
   lexer: _Lexer;
@@ -33,6 +34,11 @@ export interface RendererExtension {
 
 export type TokenizerAndRendererExtension = TokenizerExtension | RendererExtension | (TokenizerExtension & RendererExtension);
 
+type HooksApi = Omit<_Hooks, 'constructor' | 'options'>;
+type HooksObject = {
+  [K in keyof HooksApi]?: (...args: Parameters<HooksApi[K]>) => ReturnType<HooksApi[K]> | Promise<ReturnType<HooksApi[K]>>
+};
+
 type RendererApi = Omit<_Renderer, 'constructor' | 'options'>;
 type RendererObject = {
   [K in keyof RendererApi]?: (...args: Parameters<RendererApi[K]>) => ReturnType<RendererApi[K]> | false
@@ -69,14 +75,10 @@ export interface MarkedExtension {
   /**
    * Hooks are methods that hook into some part of marked.
    * preprocess is called to process markdown before sending it to marked.
+   * processAllTokens is called with the TokensList before walkTokens.
    * postprocess is called to process html after marked has finished parsing.
    */
-  hooks?: {
-    preprocess: (markdown: string) => string | Promise<string>,
-    postprocess: (html: string) => string | Promise<string>,
-    // eslint-disable-next-line no-use-before-define
-    options?: MarkedOptions
-  } | null;
+  hooks?: HooksObject | undefined | null;
 
   /**
    * Conform to obscure parts of markdown.pl as much as possible. Don't fix any of the original markdown bugs or poor behavior.
@@ -109,7 +111,12 @@ export interface MarkedExtension {
   walkTokens?: ((token: Token) => void | Promise<void>) | undefined | null;
 }
 
-export interface MarkedOptions extends Omit<MarkedExtension, 'renderer' | 'tokenizer' | 'extensions' | 'walkTokens'> {
+export interface MarkedOptions extends Omit<MarkedExtension, 'hooks' | 'renderer' | 'tokenizer' | 'extensions' | 'walkTokens'> {
+  /**
+   * Hooks are methods that hook into some part of marked.
+   */
+  hooks?: _Hooks | undefined | null;
+
   /**
    * Type: object Default: new Renderer()
    *
diff --git a/test/types/marked.ts b/test/types/marked.ts
index 571edaf789..8bf2997b2e 100644
--- a/test/types/marked.ts
+++ b/test/types/marked.ts
@@ -323,3 +323,24 @@ marked.use({
     }
   }
 });
+marked.use({
+  hooks: {
+    processAllTokens(tokens) {
+      return tokens;
+    }
+  }
+});
+marked.use({
+  async: true,
+  hooks: {
+    async preprocess(markdown) {
+      return markdown;
+    },
+    async postprocess(html) {
+      return html;
+    },
+    async processAllTokens(tokens) {
+      return tokens;
+    }
+  }
+});
diff --git a/test/unit/Hooks.test.js b/test/unit/Hooks.test.js
index cfa8ac365f..f3dfaf9035 100644
--- a/test/unit/Hooks.test.js
+++ b/test/unit/Hooks.test.js
@@ -3,6 +3,18 @@ import { timeout } from './utils.js';
 import { describe, it, beforeEach } from 'node:test';
 import assert from 'node:assert';
 
+function createHeadingToken(text) {
+  return {
+    type: 'heading',
+    raw: `# ${text}`,
+    depth: 1,
+    text,
+    tokens: [
+      { type: 'text', raw: text, text }
+    ]
+  };
+}
+
 describe('Hooks', () => {
   let marked;
   beforeEach(() => {
@@ -93,6 +105,48 @@ describe('Hooks', () => {
     assert.strictEqual(html.trim(), '<p><em>text</em></p>\n<h1>postprocess async</h1>');
   });
 
+  it('should process tokens before walkTokens', () => {
+    marked.use({
+      hooks: {
+        processAllTokens(tokens) {
+          tokens.push(createHeadingToken('processAllTokens'));
+          return tokens;
+        }
+      },
+      walkTokens(token) {
+        if (token.type === 'heading') {
+          token.tokens[0].text += ' walked';
+        }
+        return token;
+      }
+    });
+    const html = marked.parse('*text*');
+    assert.strictEqual(html.trim(), '<p><em>text</em></p>\n<h1>processAllTokens walked</h1>');
+  });
+
+  it('should process tokens async before walkTokens', async() => {
+    marked.use({
+      async: true,
+      hooks: {
+        async processAllTokens(tokens) {
+          await timeout();
+          tokens.push(createHeadingToken('processAllTokens async'));
+          return tokens;
+        }
+      },
+      walkTokens(token) {
+        if (token.type === 'heading') {
+          token.tokens[0].text += ' walked';
+        }
+        return token;
+      }
+    });
+    const promise = marked.parse('*text*');
+    assert.ok(promise instanceof Promise);
+    const html = await promise;
+    assert.strictEqual(html.trim(), '<p><em>text</em></p>\n<h1>processAllTokens async walked</h1>');
+  });
+
   it('should process all hooks in reverse', async() => {
     marked.use({
       hooks: {
@@ -101,6 +155,10 @@ describe('Hooks', () => {
         },
         postprocess(html) {
           return html + '<h1>postprocess1</h1>\n';
+        },
+        processAllTokens(tokens) {
+          tokens.push(createHeadingToken('processAllTokens1'));
+          return tokens;
         }
       }
     });
@@ -113,12 +171,23 @@ describe('Hooks', () => {
         async postprocess(html) {
           await timeout();
           return html + '<h1>postprocess2 async</h1>\n';
+        },
+        processAllTokens(tokens) {
+          tokens.push(createHeadingToken('processAllTokens2'));
+          return tokens;
         }
       }
     });
     const promise = marked.parse('*text*');
     assert.ok(promise instanceof Promise);
     const html = await promise;
-    assert.strictEqual(html.trim(), '<h1>preprocess1</h1>\n<h1>preprocess2</h1>\n<p><em>text</em></p>\n<h1>postprocess2 async</h1>\n<h1>postprocess1</h1>');
+    assert.strictEqual(html.trim(), `\
+<h1>preprocess1</h1>
+<h1>preprocess2</h1>
+<p><em>text</em></p>
+<h1>processAllTokens2</h1>
+<h1>processAllTokens1</h1>
+<h1>postprocess2 async</h1>
+<h1>postprocess1</h1>`);
   });
 });