From 4e870ae12fb7728b79705aca676d59fd2edc112c Mon Sep 17 00:00:00 2001
From: Mathias Millet <mathias@mmill.eu>
Date: Sat, 19 Oct 2024 12:07:28 +0200
Subject: [PATCH] feat: add NumberedHeadingsPreprocessor

---
 docs/source/api/preprocessors.rst             |  2 +
 nbconvert/preprocessors/__init__.py           |  2 +
 nbconvert/preprocessors/numbered_headings.py  | 43 ++++++++++
 tests/preprocessors/test_numbered_headings.py | 86 +++++++++++++++++++
 4 files changed, 133 insertions(+)
 create mode 100644 nbconvert/preprocessors/numbered_headings.py
 create mode 100644 tests/preprocessors/test_numbered_headings.py

diff --git a/docs/source/api/preprocessors.rst b/docs/source/api/preprocessors.rst
index 6276007a6..b5f74a9bc 100644
--- a/docs/source/api/preprocessors.rst
+++ b/docs/source/api/preprocessors.rst
@@ -36,6 +36,8 @@ Converting text
 
 .. autoclass:: HighlightMagicsPreprocessor
 
+.. autoclass:: NumberedHeadingsPreprocessor
+
 Metadata and header control
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/nbconvert/preprocessors/__init__.py b/nbconvert/preprocessors/__init__.py
index d752edf90..de527d0ec 100644
--- a/nbconvert/preprocessors/__init__.py
+++ b/nbconvert/preprocessors/__init__.py
@@ -13,6 +13,7 @@
 from .extractoutput import ExtractOutputPreprocessor
 from .highlightmagics import HighlightMagicsPreprocessor
 from .latex import LatexPreprocessor
+from .numbered_headings import NumberedHeadingsPreprocessor
 from .regexremove import RegexRemovePreprocessor
 from .svg2pdf import SVG2PDFPreprocessor
 from .tagremove import TagRemovePreprocessor
@@ -30,6 +31,7 @@
     "ExtractOutputPreprocessor",
     "HighlightMagicsPreprocessor",
     "LatexPreprocessor",
+    "NumberedHeadingsPreprocessor",
     "RegexRemovePreprocessor",
     "SVG2PDFPreprocessor",
     "TagRemovePreprocessor",
diff --git a/nbconvert/preprocessors/numbered_headings.py b/nbconvert/preprocessors/numbered_headings.py
new file mode 100644
index 000000000..f03235fee
--- /dev/null
+++ b/nbconvert/preprocessors/numbered_headings.py
@@ -0,0 +1,43 @@
+import re
+
+from nbconvert.preprocessors.base import Preprocessor
+
+
+class NumberedHeadingsPreprocessor(Preprocessor):
+    """Pre-processor that will rewrite markdown headings to include numberings."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.current_numbering = [0]
+
+    def format_numbering(self):
+        """Return a string representation of the current numbering"""
+        return ".".join(str(n) for n in self.current_numbering)
+
+    def inc_current_numbering(self, level):
+        if level > len(self.current_numbering):
+            self.current_numbering = self.current_numbering + [0] * (
+                level - len(self.current_numbering)
+            )
+        elif level < len(self.current_numbering):
+            self.current_numbering = self.current_numbering[:level]
+        self.current_numbering[level - 1] += 1
+
+    def transform_markdown_line(self, line, resources):
+        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):
+        if cell["cell_type"] == "markdown":
+            cell["source"] = "\n".join(
+                self.transform_markdown_line(line, resources)
+                for line in cell["source"].splitlines()
+            )
+
+        return cell, resources
diff --git a/tests/preprocessors/test_numbered_headings.py b/tests/preprocessors/test_numbered_headings.py
new file mode 100644
index 000000000..d7740f779
--- /dev/null
+++ b/tests/preprocessors/test_numbered_headings.py
@@ -0,0 +1,86 @@
+"""
+Module with tests for the Numbered Headings preprocessor.
+"""
+
+from nbformat import v4 as nbformat
+
+from nbconvert.preprocessors.numbered_headings import NumberedHeadingsPreprocessor
+
+from .base import PreprocessorTestsBase
+
+MARKDOWN_1 = """
+# Heading 1
+
+## Sub-heading
+
+some content
+"""
+
+MARKDOWN_1_POST = """
+# 1 Heading 1
+
+## 1.1 Sub-heading
+
+some content
+"""
+
+
+MARKDOWN_2 = """
+
+## Second sub-heading
+
+# Another main heading
+
+## Sub-heading
+
+
+some more content
+
+### Third heading
+"""
+
+MARKDOWN_2_POST = """
+
+## 1.2 Second sub-heading
+
+# 2 Another main heading
+
+## 2.1 Sub-heading
+
+
+some more content
+
+### 2.1.1 Third heading
+"""
+
+
+class TestNumberedHeadings(PreprocessorTestsBase):
+    def build_notebook(self):
+        cells = [
+            nbformat.new_code_cell(source="$ e $", execution_count=1),
+            nbformat.new_markdown_cell(source=MARKDOWN_1),
+            nbformat.new_code_cell(source="$ e $", execution_count=1),
+            nbformat.new_markdown_cell(source=MARKDOWN_2),
+        ]
+
+        return nbformat.new_notebook(cells=cells)
+
+    def build_preprocessor(self):
+        """Make an instance of a preprocessor"""
+        preprocessor = NumberedHeadingsPreprocessor()
+        preprocessor.enabled = True
+        return preprocessor
+
+    def test_constructor(self):
+        """Can a ClearOutputPreprocessor be constructed?"""
+        self.build_preprocessor()
+
+    def test_output(self):
+        """Test the output of the NumberedHeadingsPreprocessor"""
+        nb = self.build_notebook()
+        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()