From fe4485c55fc8167a43103985ffaf145aa8e4a08f Mon Sep 17 00:00:00 2001
From: Aakansha Doshi <aakansha1216@gmail.com>
Date: Tue, 28 Nov 2023 09:01:54 +0530
Subject: [PATCH] update mermaidAPI to cleanup the text before passing to
 getDiagramFromText

---
 packages/mermaid/src/Diagram.ts         |  4 +---
 packages/mermaid/src/diagram.spec.ts    | 10 ----------
 packages/mermaid/src/mermaidAPI.spec.ts | 13 +++++++++++++
 packages/mermaid/src/mermaidAPI.ts      |  9 +++++++--
 4 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/packages/mermaid/src/Diagram.ts b/packages/mermaid/src/Diagram.ts
index 5f17d09927..2b5ae89923 100644
--- a/packages/mermaid/src/Diagram.ts
+++ b/packages/mermaid/src/Diagram.ts
@@ -8,8 +8,6 @@ import { encodeEntities } from './utils.js';
 import type { DetailedError } from './utils.js';
 import type { DiagramDefinition, DiagramMetadata } from './diagram-api/types.js';
 
-import { cleanupComments } from './diagram-api/comments.js';
-
 export type ParseErrorFunction = (err: string | DetailedError | unknown, hash?: any) => void;
 
 /**
@@ -25,7 +23,7 @@ export class Diagram {
 
   private detectError?: UnknownDiagramError;
   constructor(public text: string, public metadata: Pick<DiagramMetadata, 'title'> = {}) {
-    this.text = encodeEntities(cleanupComments(text));
+    this.text = encodeEntities(text);
     this.text += '\n';
     const cnf = configApi.getConfig();
     try {
diff --git a/packages/mermaid/src/diagram.spec.ts b/packages/mermaid/src/diagram.spec.ts
index 7a3daffbee..4354f57bce 100644
--- a/packages/mermaid/src/diagram.spec.ts
+++ b/packages/mermaid/src/diagram.spec.ts
@@ -83,14 +83,4 @@ Expecting 'TXT', got 'NEWLINE'"
     expect(messages[0].message).toBe('I fl°°9829¶ß you!');
     expect(messages[1].message).toBe('I fl°°9829¶ß you fl°infin¶ß times more!');
   });
-
-  test('should clean up comments inside getDiagramFromText when present in diagram definition', async () => {
-    const diagram = await getDiagramFromText(
-      `flowchart LR
-    %% this is a comment A -- text --> B{node}
-    A -- text --> B -- text2 --> C`
-    );
-    expect(diagram).toBeInstanceOf(Diagram);
-    expect(diagram.type).toBe('flowchart-v2');
-  });
 });
diff --git a/packages/mermaid/src/mermaidAPI.spec.ts b/packages/mermaid/src/mermaidAPI.spec.ts
index 32647ee625..7e8f46d06a 100644
--- a/packages/mermaid/src/mermaidAPI.spec.ts
+++ b/packages/mermaid/src/mermaidAPI.spec.ts
@@ -67,6 +67,7 @@ vi.mock('stylis', () => {
 });
 import { compile, serialize } from 'stylis';
 import { decodeEntities, encodeEntities } from './utils.js';
+import { Diagram } from './Diagram.js';
 
 /**
  * @see https://vitest.dev/guide/mocking.html Mock part of a module
@@ -744,4 +745,16 @@ describe('mermaidAPI', () => {
       });
     });
   });
+
+  describe('getDiagramFromText', () => {
+    it('should clean up comments when present in diagram definition', async () => {
+      const diagram = await mermaidAPI.getDiagramFromText(
+        `flowchart LR
+      %% this is a comment A -- text --> B{node}
+      A -- text --> B -- text2 --> C`
+      );
+      expect(diagram).toBeInstanceOf(Diagram);
+      expect(diagram.type).toBe('flowchart-v2');
+    });
+  });
 });
diff --git a/packages/mermaid/src/mermaidAPI.ts b/packages/mermaid/src/mermaidAPI.ts
index 09d0077803..65310d4fae 100644
--- a/packages/mermaid/src/mermaidAPI.ts
+++ b/packages/mermaid/src/mermaidAPI.ts
@@ -17,7 +17,7 @@ import { compile, serialize, stringify } from 'stylis';
 import { version } from '../package.json';
 import * as configApi from './config.js';
 import { addDiagrams } from './diagram-api/diagram-orchestration.js';
-import { Diagram, getDiagramFromText } from './Diagram.js';
+import { Diagram, getDiagramFromText as getDiagramFromTextInternal } from './Diagram.js';
 import errorRenderer from './diagrams/error/errorRenderer.js';
 import { attachFunctions } from './interactionDb.js';
 import { log, setLogLevel } from './logger.js';
@@ -28,7 +28,7 @@ import type { MermaidConfig } from './config.type.js';
 import { evaluate } from './diagrams/common/common.js';
 import isEmpty from 'lodash-es/isEmpty.js';
 import { setA11yDiagramInfo, addSVGa11yTitleDescription } from './accessibility.js';
-import type { DiagramStyleClassDef } from './diagram-api/types.js';
+import type { DiagramMetadata, DiagramStyleClassDef } from './diagram-api/types.js';
 import { preprocessDiagram } from './preprocess.js';
 import { decodeEntities } from './utils.js';
 
@@ -519,6 +519,11 @@ function initialize(options: MermaidConfig = {}) {
   addDiagrams();
 }
 
+const getDiagramFromText = (text: string, metadata: Pick<DiagramMetadata, 'title'> = {}) => {
+  const { code } = preprocessDiagram(text);
+  return getDiagramFromTextInternal(code, metadata);
+};
+
 /**
  * Add accessibility (a11y) information to the diagram.
  *