Merge pull request #21 from Mermaid-Chart/fix/keep-comments-in-yaml

fix(cli): keep comments in YAML frontmatter

Author: aloisklink

packages/cli/package.json

@@ -59,11 +59,11 @@
     "@inquirer/select": "^1.3.1",
     "@mermaidchart/sdk": "workspace:^",
     "commander": "^11.1.0",
-    "js-yaml": "^4.1.0",
     "remark": "^15.0.1",
     "remark-frontmatter": "^5.0.0",
     "remark-gfm": "^4.0.0",
     "to-vfile": "^8.0.0",
-    "unist-util-visit": "^5.0.0"
+    "unist-util-visit": "^5.0.0",
+    "yaml": "^2.3.4"
   }
 }

packages/cli/src/commander.test.ts

@@ -435,6 +435,7 @@ describe('pull', () => {
     const mockedDiagram = {
       ...mockedEmptyDiagram,
       code: `---
+# comments in YAML shouldn't be removed
 title: My cool flowchart
 ---
   flowchart TD
@@ -453,6 +454,7 @@ title: My cool flowchart
       expect(diagramContents).toContain(
         `id: https://test.mermaidchart.invalid/d/${mockedDiagram.documentID}`,
       );
+      expect(diagramContents).toContain("# comments in YAML shouldn't be removed");
       expect(diagramContents).toContain("flowchart TD\n      A[I've been updated!]");
     }
   });
@@ -521,6 +523,11 @@ describe('push', () => {
         code: expect.not.stringContaining('id:'),
       }),
     );
+    expect(vi.mocked(MermaidChart.prototype.setDocument)).toHaveBeenCalledWith(
+      expect.objectContaining({
+        code: expect.stringMatching(/^# This comment should be pushed to the server/m),
+      }),
+    );
   });
 
   it('should push documents from within markdown file', async () => {

packages/cli/src/frontmatter.ts

@@ -1,9 +1,7 @@
 /**
  * Copied from https://github.com/mermaid-js/mermaid/blob/4a4e614b646bdb5f91f02d0483a7704b315d09fd/packages/mermaid/src/diagram-api/regexes.ts
  */
-
-// The "* as yaml" part is necessary for tree-shaking
-import * as yaml from 'js-yaml';
+import { parseDocument, type Document, YAMLMap, isMap } from 'yaml';
 
 const frontMatterRegex = /^-{3}\s*[\n\r](.*?)[\n\r]-{3}\s*[\n\r]+/s;
 const urlIDRegex = /(?<baseURL>.*)\/d\/(?<documentID>[\w-]+)/;
@@ -59,20 +57,13 @@ function splitFrontMatter(text: string) {
   }
 }
 
-function parseFrontMatterYAML(frontMatterYaml: string) {
-  let parsed: FrontMatterMetadata =
-    // TODO: replace with https://www.npmjs.com/package/yaml so that we can
-    //       read/write comments too
-    yaml.load(frontMatterYaml, {
-      // To support config, we need JSON schema.
-      // https://www.yaml.org/spec/1.2/spec.html#id2803231
-      schema: yaml.JSON_SCHEMA,
-    }) ?? {};
-
-  // To handle runtime data type changes
-  parsed = typeof parsed === 'object' && !Array.isArray(parsed) ? parsed : {};
+function parseFrontMatterYAML(frontMatterYaml: string): Document<YAMLMap, false> {
+  const document: Document = parseDocument(frontMatterYaml);
+  if (!isMap(document.contents)) {
+    document.contents = new YAMLMap();
+  }
 
-  return parsed;
+  return document as unknown as Document<YAMLMap, false>;
 }
 
 /**
@@ -85,7 +76,7 @@ function parseFrontMatterYAML(frontMatterYaml: string) {
 export function extractFrontMatter(text: string): FrontMatterResult {
   const { diagramText, frontMatter } = splitFrontMatter(text);
 
-  const parsed = parseFrontMatterYAML(frontMatter);
+  const parsed = parseFrontMatterYAML(frontMatter).toJSON();
 
   const metadata: FrontMatterMetadata = {};
 
@@ -116,16 +107,16 @@ export function extractFrontMatter(text: string): FrontMatterResult {
  * @param newMetadata - The metadata fields to update.
  * @returns The text with the updated YAML frontmatter.
  */
-export function injectFrontMatter(text: string, newMetadata: Partial<FrontMatterMetadata>) {
+export function injectFrontMatter(text: string, newMetadata: Pick<FrontMatterMetadata, 'id'>) {
   const { diagramText, frontMatter } = splitFrontMatter(text);
 
-  const parsed = parseFrontMatterYAML(frontMatter);
+  const document = parseFrontMatterYAML(frontMatter);
 
-  const mergedFrontmatter = { ...parsed, ...newMetadata };
+  for (const [key, value] of Object.entries(newMetadata)) {
+    document.contents.set(key, value);
+  }
 
-  return `---\n${yaml.dump(mergedFrontmatter, {
-    schema: yaml.JSON_SCHEMA,
-  })}---\n${diagramText}`;
+  return `---\n${document.toString()}---\n${diagramText}`;
 }
 
 /**
@@ -138,24 +129,16 @@ export function injectFrontMatter(text: string, newMetadata: Partial<FrontMatter
 export function removeFrontMatterKeys(text: string, keysToRemove: Set<keyof FrontMatterMetadata>) {
   const { diagramText, frontMatter } = splitFrontMatter(text);
 
-  const parsedFrontMatter = parseFrontMatterYAML(frontMatter);
+  const document = parseFrontMatterYAML(frontMatter);
 
-  const entries = Object.entries(parsedFrontMatter)
-    .map((val) => {
-      if (keysToRemove.has(val[0] as keyof FrontMatterMetadata)) {
-        return null;
-      } else {
-        return val;
-      }
-    })
-    .filter((val) => val) as [string, any][]; // eslint-disable-line @typescript-eslint/no-explicit-any
+  for (const key of keysToRemove) {
+    document.contents.delete(key);
+  }
 
-  if (entries.length === 0) {
+  if (document.contents.items.length === 0) {
     // skip creating frontmatter if there is no frontmatter
     return diagramText;
   } else {
-    return `---\n${yaml.dump(Object.fromEntries(entries), {
-      schema: yaml.JSON_SCHEMA,
-    })}---\n${diagramText}`;
+    return `---\n${document.toString()}---\n${diagramText}`;
   }
 }

packages/cli/test/fixtures/connected-diagram.mmd

@@ -1,4 +1,5 @@
 ---
+# This comment should be pushed to the server
 title: My cool flowchart
 id: https://test.mermaidchart.invalid/d/my-test-document-id
 ---