feat: Add support for config in frontmatter

mermaid-js · Aug 21, 2023 · 85d2b8f · 85d2b8f
1 parent 767baa4
commit 85d2b8f
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 32 deletions.
diff --git a/packages/mermaid/src/Diagram.ts b/packages/mermaid/src/Diagram.ts
@@ -48,7 +48,7 @@ export class Diagram {
     // extractFrontMatter().
 
     this.parser.parse = (text: string) =>
-      originalParse(cleanupComments(extractFrontMatter(text, this.db)));
+      originalParse(cleanupComments(extractFrontMatter(text, this.db, configApi.addDirective)));
 
     this.parser.parser.yy = this.db;
     this.init = diagram.init;

diff --git a/packages/mermaid/src/diagram-api/frontmatter.spec.ts b/packages/mermaid/src/diagram-api/frontmatter.spec.ts
@@ -2,76 +2,94 @@ import { vi } from 'vitest';
 import { extractFrontMatter } from './frontmatter.js';
 
 const dbMock = () => ({ setDiagramTitle: vi.fn() });
+const setConfigMock = vi.fn();
 
 describe('extractFrontmatter', () => {
+  beforeEach(() => {
+    setConfigMock.mockClear();
+  });
+
   it('returns text unchanged if no frontmatter', () => {
-    expect(extractFrontMatter('diagram', dbMock())).toEqual('diagram');
+    expect(extractFrontMatter('diagram', dbMock(), setConfigMock)).toEqual('diagram');
   });
 
   it('returns text unchanged if frontmatter lacks closing delimiter', () => {
     const text = `---\ntitle: foo\ndiagram`;
-    expect(extractFrontMatter(text, dbMock())).toEqual(text);
+    expect(extractFrontMatter(text, dbMock(), setConfigMock)).toEqual(text);
   });
 
   it('handles empty frontmatter', () => {
     const db = dbMock();
     const text = `---\n\n---\ndiagram`;
-    expect(extractFrontMatter(text, db)).toEqual('diagram');
+    expect(extractFrontMatter(text, db, setConfigMock)).toEqual('diagram');
     expect(db.setDiagramTitle).not.toHaveBeenCalled();
   });
 
   it('handles frontmatter without mappings', () => {
     const db = dbMock();
     const text = `---\n1\n---\ndiagram`;
-    expect(extractFrontMatter(text, db)).toEqual('diagram');
+    expect(extractFrontMatter(text, db, setConfigMock)).toEqual('diagram');
     expect(db.setDiagramTitle).not.toHaveBeenCalled();
   });
 
   it('does not try to parse frontmatter at the end', () => {
     const db = dbMock();
     const text = `diagram\n---\ntitle: foo\n---\n`;
-    expect(extractFrontMatter(text, db)).toEqual(text);
+    expect(extractFrontMatter(text, db, setConfigMock)).toEqual(text);
     expect(db.setDiagramTitle).not.toHaveBeenCalled();
   });
 
   it('handles frontmatter with multiple delimiters', () => {
     const db = dbMock();
     const text = `---\ntitle: foo---bar\n---\ndiagram\n---\ntest`;
-    expect(extractFrontMatter(text, db)).toEqual('diagram\n---\ntest');
+    expect(extractFrontMatter(text, db, setConfigMock)).toEqual('diagram\n---\ntest');
     expect(db.setDiagramTitle).toHaveBeenCalledWith('foo---bar');
   });
 
   it('handles frontmatter with multi-line string and multiple delimiters', () => {
     const db = dbMock();
     const text = `---\ntitle: |\n   multi-line string\n   ---\n---\ndiagram`;
-    expect(extractFrontMatter(text, db)).toEqual('diagram');
+    expect(extractFrontMatter(text, db, setConfigMock)).toEqual('diagram');
     expect(db.setDiagramTitle).toHaveBeenCalledWith('multi-line string\n---\n');
   });
 
   it('handles frontmatter with title', () => {
     const db = dbMock();
     const text = `---\ntitle: foo\n---\ndiagram`;
-    expect(extractFrontMatter(text, db)).toEqual('diagram');
+    expect(extractFrontMatter(text, db, setConfigMock)).toEqual('diagram');
     expect(db.setDiagramTitle).toHaveBeenCalledWith('foo');
   });
 
+  it('handles frontmatter with config', () => {
+    const db = dbMock();
+    const text = `---
+title: hello
+config:
+  flowchart:
+    htmlLabels: false
+---
+diagram`;
+    expect(extractFrontMatter(text, db, setConfigMock)).toEqual('diagram');
+    expect(setConfigMock).toHaveBeenCalledWith({ flowchart: { htmlLabels: false } });
+  });
+
   it('handles booleans in frontmatter properly', () => {
     const db = dbMock();
     const text = `---\ntitle: true\n---\ndiagram`;
-    expect(extractFrontMatter(text, db)).toEqual('diagram');
+    expect(extractFrontMatter(text, db, setConfigMock)).toEqual('diagram');
     expect(db.setDiagramTitle).toHaveBeenCalledWith('true');
   });
 
   it('ignores unspecified frontmatter keys', () => {
     const db = dbMock();
     const text = `---\ninvalid: true\ntitle: foo\ntest: bar\n---\ndiagram`;
-    expect(extractFrontMatter(text, db)).toEqual('diagram');
+    expect(extractFrontMatter(text, db, setConfigMock)).toEqual('diagram');
     expect(db.setDiagramTitle).toHaveBeenCalledWith('foo');
   });
 
   it('throws exception for invalid YAML syntax', () => {
     const text = `---\n!!!\n---\ndiagram`;
-    expect(() => extractFrontMatter(text, dbMock())).toThrow(
+    expect(() => extractFrontMatter(text, dbMock(), setConfigMock)).toThrow(
       'tag suffix cannot contain exclamation marks'
     );
   });

diff --git a/packages/mermaid/src/diagram-api/frontmatter.ts b/packages/mermaid/src/diagram-api/frontmatter.ts
@@ -1,3 +1,4 @@
+import { MermaidConfig } from '../config.type.js';
 import { DiagramDB } from './types.js';
 // The "* as yaml" part is necessary for tree-shaking
 import * as yaml from 'js-yaml';
@@ -9,11 +10,12 @@ import * as yaml from 'js-yaml';
 // Relevant YAML spec: https://yaml.org/spec/1.2.2/#914-explicit-documents
 export const frontMatterRegex = /^-{3}\s*[\n\r](.*?)[\n\r]-{3}\s*[\n\r]+/s;
 
-type FrontMatterMetadata = {
+interface FrontMatterMetadata {
   title?: string;
   // Allows custom display modes. Currently used for compact mode in gantt charts.
   displayMode?: string;
-};
+  config?: MermaidConfig;
+}
 
 /**
  * Extract and parse frontmatter from text, if present, and sets appropriate
@@ -22,25 +24,35 @@ type FrontMatterMetadata = {
  * @param db - Diagram database, could be of any diagram.
  * @returns text with frontmatter stripped out
  */
-export function extractFrontMatter(text: string, db: DiagramDB): string {
+export function extractFrontMatter(
+  text: string,
+  db: DiagramDB,
+  setDiagramConfig: (config: MermaidConfig) => void
+): string {
   const matches = text.match(frontMatterRegex);
-  if (matches) {
-    const parsed: FrontMatterMetadata = yaml.load(matches[1], {
-      // To keep things simple, only allow strings, arrays, and plain objects.
-      // https://www.yaml.org/spec/1.2/spec.html#id2802346
-      schema: yaml.FAILSAFE_SCHEMA,
-    }) as FrontMatterMetadata;
-
-    if (parsed?.title) {
-      db.setDiagramTitle?.(parsed.title);
-    }
-
-    if (parsed?.displayMode) {
-      db.setDisplayMode?.(parsed.displayMode);
-    }
-
-    return text.slice(matches[0].length);
-  } else {
+  if (!matches) {
     return text;
   }
+
+  const parsed: FrontMatterMetadata = yaml.load(matches[1], {
+    // To support config, we need JSON schema.
+    // https://www.yaml.org/spec/1.2/spec.html#id2803231
+    schema: yaml.JSON_SCHEMA,
+  }) as FrontMatterMetadata;
+
+  if (parsed?.title) {
+    // toString() is necessary because YAML could parse the title as a number/boolean
+    db.setDiagramTitle?.(parsed.title.toString());
+  }
+
+  if (parsed?.displayMode) {
+    // toString() is necessary because YAML could parse the title as a number/boolean
+    db.setDisplayMode?.(parsed.displayMode.toString());
+  }
+
+  if (parsed?.config) {
+    setDiagramConfig(parsed.config);
+  }
+
+  return text.slice(matches[0].length);
 }