From 20cb7f5fdb145b13967d61fe78363c6db5a287f8 Mon Sep 17 00:00:00 2001
From: Daniel Kuznetsov <letux.dev@gmail.com>
Date: Tue, 10 Dec 2024 18:25:19 +0300
Subject: [PATCH] feat: implement ways to interact with collect via transform
 API

---
 src/transform/index.ts   | 47 +++++++++++++++++++++++++++++++++++-----
 src/transform/typings.ts | 25 +++++++++++++++++++--
 2 files changed, 64 insertions(+), 8 deletions(-)

diff --git a/src/transform/index.ts b/src/transform/index.ts
index 3ca9dfa3..6588de75 100644
--- a/src/transform/index.ts
+++ b/src/transform/index.ts
@@ -1,9 +1,9 @@
-import type {EnvType, OptionsType, OutputType} from './typings';
+import type {EnvType, OptionsType, OutputType, RootCollectorOptions} from './typings';
 
 import {bold} from 'chalk';
 
 import {log} from './log';
-import liquidSnippet from './liquid';
+import liquidDocument from './liquid';
 import initMarkdownIt from './md';
 
 function applyLiquid(input: string, options: OptionsType) {
@@ -17,7 +17,7 @@ function applyLiquid(input: string, options: OptionsType) {
 
     return disableLiquid || isLiquided
         ? input
-        : liquidSnippet(input, vars, path, {conditionsInCode});
+        : liquidDocument(input, vars, path, {conditionsInCode});
 }
 
 function handleError(error: unknown, path?: string): never {
@@ -33,8 +33,14 @@ function emitResult(html: string, env: EnvType): OutputType {
     };
 }
 
+type TransformFunction = {
+    (originInput: string, options?: OptionsType): OutputType;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    collect: (input: string, options: RootCollectorOptions<any>) => string;
+};
+
 // eslint-disable-next-line consistent-return
-function transform(originInput: string, options: OptionsType = {}) {
+const transform: TransformFunction = (originInput: string, options: OptionsType = {}) => {
     const input = applyLiquid(originInput, options);
     const {parse, compile, env} = initMarkdownIt(options);
 
@@ -43,11 +49,40 @@ function transform(originInput: string, options: OptionsType = {}) {
     } catch (error) {
         handleError(error, options.path);
     }
-}
+};
+
+transform.collect = (
+    input: string,
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    {mdItInitOptions, pluginCollectOptions, parserPluginsOverride}: RootCollectorOptions<any>,
+) => {
+    const maybeLiquidedInput = applyLiquid(input, mdItInitOptions);
+    const {parse} = initMarkdownIt({
+        ...mdItInitOptions,
+        plugins: parserPluginsOverride ?? mdItInitOptions.plugins,
+    });
+
+    const plugins = mdItInitOptions.plugins ?? [];
+
+    try {
+        const tokenStream = parse(maybeLiquidedInput);
+
+        return plugins.reduce((collected, plugin) => {
+            const collectOutput = plugin.collect?.(collected, {
+                ...pluginCollectOptions,
+                tokenStream,
+            });
+
+            return collectOutput ?? collected;
+        }, input);
+    } catch (error) {
+        handleError(error, mdItInitOptions.path);
+    }
+};
 
 export = transform;
 
-// eslint-disable-next-line @typescript-eslint/no-namespace, no-redeclare -- backward compatibility
+// eslint-disable-next-line @typescript-eslint/no-namespace, @typescript-eslint/no-redeclare -- backward compatibility
 namespace transform {
     export type Options = OptionsType;
     export type Output = OutputType;
diff --git a/src/transform/typings.ts b/src/transform/typings.ts
index 0b67a801..cd5210d1 100644
--- a/src/transform/typings.ts
+++ b/src/transform/typings.ts
@@ -1,5 +1,5 @@
 import {LanguageFn} from 'highlight.js';
-import DefaultMarkdownIt from 'markdown-it';
+import DefaultMarkdownIt, {Token} from 'markdown-it';
 import DefaultStateCore from 'markdown-it/lib/rules_core/state_core';
 
 import {SanitizeOptions} from './sanitize';
@@ -47,7 +47,7 @@ export interface OptionsType {
     sanitizeOptions?: SanitizeOptions;
     needFlatListHeadings?: boolean;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    plugins?: MarkdownItPluginCb<any>[];
+    plugins?: ExtendedPluginWithCollect<any, any>[];
     preprocessors?: MarkdownItPreprocessorCb[]; // Preprocessors should modify the input before passing it to MD
     highlightLangs?: HighlightLangMap;
     disableRules?: string[];
@@ -96,6 +96,27 @@ export type MarkdownItPluginCb<T extends {} = {}> = {
     (md: MarkdownIt, opts: T & MarkdownItPluginOpts): void;
 };
 
+export type IntrinsicCollectOptions = {
+    tokenStream: Token[];
+};
+
+export type ExtendedPluginWithCollect<
+    PluginRegularOptions extends {} = {},
+    PluginCollectOptions = {},
+> = MarkdownItPluginCb<PluginRegularOptions> & {
+    collect?: (
+        input: string,
+        options: PluginCollectOptions & IntrinsicCollectOptions,
+    ) => string | void;
+};
+
+export type RootCollectorOptions<PluginCollectOptions> = {
+    mdItInitOptions: OptionsType;
+    pluginCollectOptions: PluginCollectOptions;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    parserPluginsOverride?: ExtendedPluginWithCollect<any, any>[];
+};
+
 export type MarkdownItPreprocessorCb<T extends unknown = {}> = {
     (input: string, opts: T & Partial<MarkdownItPluginOpts>, md?: MarkdownIt): string;
 };