Does frontmatter check before loop, adds test case

jupyterlab · Mar 17, 2023 · 1fca3f8 · 1fca3f8
1 parent 16c5e36
commit 1fca3f8
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 17 deletions.
diff --git a/packages/toc/src/utils/markdown.ts b/packages/toc/src/utils/markdown.ts
@@ -65,8 +65,27 @@ export function getHeadings(text: string): IMarkdownHeading[] {
   // Iterate over the lines to get the header level and text for each line:
   const headings = new Array<IMarkdownHeading>();
   let isCodeBlock;
-  let isFrontmatterBlock;
-  for (let lineIdx = 0; lineIdx < lines.length; lineIdx++) {
+  let lineIdx = 0;
+
+  // Don't check for Markdown headings if in a YAML frontmatter block.
+  // We can only start a frontmatter block on the first line of the file.
+  // At other positions in a markdown file, '---' represents a horizontal rule.
+  if (lines[lineIdx] === '---' ) {
+    // Search for another '---' and treat that as the end of the frontmatter.
+    // If we don't find one, treat the file as containing no frontmatter.
+    for (
+      let frontmatterEndLineIdx = lineIdx + 1;
+      frontmatterEndLineIdx < lines.length;
+      frontmatterEndLineIdx++
+    ) {
+      if (lines[frontmatterEndLineIdx] === '---') {
+        lineIdx = frontmatterEndLineIdx + 1;
+        break;
+      }
+    }
+  }
+
+  for (; lineIdx < lines.length; lineIdx++) {
     const line = lines[lineIdx];
 
     if (line === '') {
@@ -82,17 +101,6 @@ export function getHeadings(text: string): IMarkdownHeading[] {
       continue;
     }
 
-    // Don't check for Markdown headings if in a YAML frontmatter block.
-    // We can only start a frontmatter block on the first line of the file.
-    // At other positions in a markdown file, '---' represents a horizontal rule.
-    if (line === '---' && (isFrontmatterBlock || lineIdx === 0)) {
-      isFrontmatterBlock = !isFrontmatterBlock;
-    }
-
-    if (isFrontmatterBlock) {
-      continue;
-    }
-
     const heading = parseHeading(line, lines[lineIdx + 1]); // append the next line to capture alternative style Markdown headings
 
     if (heading) {

diff --git a/packages/toc/test/markdown.spec.ts b/packages/toc/test/markdown.spec.ts
@@ -230,7 +230,8 @@ front: matter
               skip: false
             }
           ]
-        ][
+        ],
+        [
           `---
 front: matter
 ---
@@ -240,8 +241,7 @@ front: matter
 # Header between horizontal rules
 ---
 
-# Header after horizontal rules`
-        ],
+# Header after horizontal rules`,
         [
           {
             text: 'Header',
@@ -268,7 +268,21 @@ front: matter
             skip: false
           }
         ]
-      ])('should extract headings from %s', (src, headers) => {
+      ],
+      [
+        `---
+# Header`,
+        [
+          {
+            text: 'Header',
+            level: 1,
+            line: 1,
+            raw: '# Header',
+            prefix: '1. ',
+            skip: false
+          }
+        ]
+      ]])('should extract headings from %s', (src, headers) => {
         const headings = TableOfContentsUtils.filterHeadings(
           TableOfContentsUtils.Markdown.getHeadings(src),
           {