From c012c1d9f194f492dcb055638c790580aa91c51e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timoth=C3=A9e=20Mazzucotelli?= <dev@pawamoy.fr>
Date: Thu, 13 Jun 2024 01:27:03 +0200
Subject: [PATCH] fix: Fix removal of temporary div for headings forwarding

Issue-50: https://github.com/pawamoy/markdown-exec/issues/50
---
 src/markdown_exec/processors.py | 29 +++++++++++++++++------------
 tests/conftest.py               |  2 +-
 tests/test_headings.py          | 30 ++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 13 deletions(-)
 create mode 100644 tests/test_headings.py

diff --git a/src/markdown_exec/processors.py b/src/markdown_exec/processors.py
index e335730..e53a15d 100644
--- a/src/markdown_exec/processors.py
+++ b/src/markdown_exec/processors.py
@@ -92,9 +92,9 @@ def run(self, root: Element) -> None:  # noqa: D102 (ignore missing docstring)
             if match:
                 counter = int(match.group(1))
                 markup: Markup = self.md.htmlStash.rawHtmlBlocks[counter]  # type: ignore[assignment]
-                if markup in self.headings:
+                if headings := self.headings.get(markup):
                     div = Element("div", {"class": "markdown-exec"})
-                    div.extend(self.headings[markup])
+                    div.extend(headings)
                     el.append(div)
 
 
@@ -104,15 +104,20 @@ class RemoveHeadings(Treeprocessor):
     name = "markdown_exec_remove_headings"
 
     def run(self, root: Element) -> None:  # noqa: D102
+        self._remove_duplicated_headings(root)
+
+    def _remove_duplicated_headings(self, parent: Element) -> None:
         carry_text = ""
-        for el in reversed(root):  # Reversed mainly for the ability to mutate during iteration.
-            for subel in reversed(el):
-                if subel.tag == "div" and subel.get("class") == "markdown-exec":
-                    # Delete the duplicated headings along with their container, but keep the text (i.e. the actual HTML).
-                    carry_text = (subel.text or "") + carry_text
-                    el.remove(subel)
-                elif carry_text:
-                    subel.tail = (subel.tail or "") + carry_text
+        for el in reversed(parent):  # Reversed mainly for the ability to mutate during iteration.
+            if el.tag == "div" and el.get("class") == "markdown-exec":
+                # Delete the duplicated headings along with their container, but keep the text (i.e. the actual HTML).
+                carry_text = (el.text or "") + carry_text
+                parent.remove(el)
+            else:
+                if carry_text:
+                    el.tail = (el.tail or "") + carry_text
                     carry_text = ""
-            if carry_text:
-                el.text = (el.text or "") + carry_text
+                self._remove_duplicated_headings(el)
+
+        if carry_text:
+            parent.text = (parent.text or "") + carry_text
diff --git a/tests/conftest.py b/tests/conftest.py
index 1512961..e6b1857 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -23,6 +23,6 @@ def md() -> Markdown:
         for language in formatters
     ]
     return Markdown(
-        extensions=["pymdownx.superfences"],
+        extensions=["pymdownx.superfences", "pymdownx.tabbed"],
         extension_configs={"pymdownx.superfences": {"custom_fences": fences}},
     )
diff --git a/tests/test_headings.py b/tests/test_headings.py
new file mode 100644
index 0000000..2b412ab
--- /dev/null
+++ b/tests/test_headings.py
@@ -0,0 +1,30 @@
+"""Tests for headings."""
+
+from __future__ import annotations
+
+from textwrap import dedent
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from markdown import Markdown
+
+
+def test_headings_removal(md: Markdown) -> None:
+    """Headings should leave no trace behind.
+
+    Parameters:
+        md: A Markdown instance (fixture).
+    """
+    html = md.convert(
+        dedent(
+            """
+            === "File layout"
+
+                ```tree
+                ./
+                    hello.md
+                ```
+            """,
+        ),
+    )
+    assert 'class="markdown-exec"' not in html