use mistune for numbered headings

Author: Mathias Millet

nbconvert/preprocessors/numbered_headings.py

@@ -2,17 +2,33 @@
 Preprocessor that transforms markdown cells: Insert numbering in from of heading
 """
 
-import re
+from traitlets.log import get_logger
 
 from nbconvert.preprocessors.base import Preprocessor
 
+logger = get_logger()
+
+try:  # for Mistune >= 3.0
+    import mistune
+    from mistune.core import BlockState
+    from mistune.renderers.markdown import MarkdownRenderer
+
+    MISTUNE_V3 = True
+except ImportError:  # for Mistune >= 2.0
+    MISTUNE_V3 = False
+
 
 class NumberedHeadingsPreprocessor(Preprocessor):
     """Pre-processor that will rewrite markdown headings to include numberings."""
 
     def __init__(self, *args, **kwargs):
         """Init"""
         super().__init__(*args, **kwargs)
+        if not MISTUNE_V3:
+            logger.error("NumberedHeadingsPreprocessor requires mistune >= 3")
+            return
+        self.md_parser = mistune.create_markdown(renderer=None)
+        self.md_renderer = MarkdownRenderer()
         self.current_numbering = [0]
 
     def format_numbering(self):
@@ -29,23 +45,26 @@ def _inc_current_numbering(self, level):
             self.current_numbering = self.current_numbering[:level]
         self.current_numbering[level - 1] += 1
 
-    def _transform_markdown_line(self, line, resources):
-        """Rewrites one markdown line, if needed"""
-        if m := re.match(r"^(?P<level>#+) (?P<heading>.*)", line):
-            level = len(m.group("level"))
-            self._inc_current_numbering(level)
-            old_heading = m.group("heading").strip()
-            new_heading = self.format_numbering() + " " + old_heading
-            return "#" * level + " " + new_heading
-
-        return line
-
     def preprocess_cell(self, cell, resources, index):
         """Rewrites all the headings in the cell if it is markdown"""
-        if cell["cell_type"] == "markdown":
-            cell["source"] = "\n".join(
-                self._transform_markdown_line(line, resources)
-                for line in cell["source"].splitlines()
-            )
-
-        return cell, resources
+        if not MISTUNE_V3:
+            return cell, resources
+        if cell["cell_type"] != "markdown":
+            return cell, resources
+        try:
+            md_ast = self.md_parser(cell["source"])
+            assert not isinstance(md_ast, str)  # type guard ; str is not returned by ast parser
+            for element in md_ast:
+                if element["type"] == "heading":
+                    level = element["attrs"]["level"]
+                    self._inc_current_numbering(level)
+                    if len(element["children"]) > 0:
+                        child = element["children"][0]
+                        if child["type"] == "text":
+                            child["raw"] = self.format_numbering() + " " + child["raw"]
+            new_source = self.md_renderer(md_ast, BlockState())
+            cell["source"] = new_source
+            return cell, resources
+        except Exception:
+            logger.warning("Failed processing cell headings", exc_info=True)
+            return cell, resources

tests/preprocessors/test_numbered_headings.py

@@ -47,12 +47,31 @@
 
 ## 2.1 Sub-heading
 
-
 some more content
 
 ### 2.1.1 Third heading
 """
 
+MARKDOWN_3 = """
+# HEADING
+
+```
+# this is not a heading
+
+## this neither
+```
+"""
+
+MARKDOWN_3_POST = """
+# 3 HEADING
+
+```
+# this is not a heading
+
+## this neither
+```
+"""
+
 
 class TestNumberedHeadings(PreprocessorTestsBase):
     def build_notebook(self):
@@ -61,6 +80,7 @@ def build_notebook(self):
             nbformat.new_markdown_cell(source=MARKDOWN_1),
             nbformat.new_code_cell(source="$ e $", execution_count=1),
             nbformat.new_markdown_cell(source=MARKDOWN_2),
+            nbformat.new_markdown_cell(source=MARKDOWN_3),
         ]
 
         return nbformat.new_notebook(cells=cells)
@@ -72,7 +92,7 @@ def build_preprocessor(self):
         return preprocessor
 
     def test_constructor(self):
-        """Can a ClearOutputPreprocessor be constructed?"""
+        """Can a NumberedHeadingsPreprocessor be constructed?"""
         self.build_preprocessor()
 
     def test_output(self):
@@ -81,6 +101,6 @@ def test_output(self):
         res = self.build_resources()
         preprocessor = self.build_preprocessor()
         nb, res = preprocessor(nb, res)
-        print(nb.cells[1].source)
         assert nb.cells[1].source.strip() == MARKDOWN_1_POST.strip()
         assert nb.cells[3].source.strip() == MARKDOWN_2_POST.strip()
+        assert nb.cells[4].source.strip() == MARKDOWN_3_POST.strip()