From 064bc8fc84e0b43853b2836db36eaac9c0a5eae6 Mon Sep 17 00:00:00 2001
From: Daniel Fichtinger <daniel@ficd.ca>
Date: Sat, 5 Jul 2025 01:58:54 -0400
Subject: [PATCH] implement image caption processing

---
 src/zona/markdown.py | 48 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/zona/markdown.py b/src/zona/markdown.py
index c28a10b..8bfc253 100644
--- a/src/zona/markdown.py
+++ b/src/zona/markdown.py
@@ -29,7 +29,46 @@ from zona.log import get_logger
 logger = get_logger()
 
 
-class ZonaLinkProcessor(Treeprocessor):
+class ZonaImageTreeprocessor(Treeprocessor):
+    """Implement Zona's image caption rendering."""
+
+    def __init__(self, md: Markdown):
+        super().__init__()
+        self.md: Markdown = md
+
+    @override
+    def run(self, root: etree.Element):
+        for parent in root.iter():
+            for idx, child in enumerate(list(parent)):
+                if (
+                    child.tag == "p"
+                    and len(child) == 1
+                    and child[0].tag == "img"
+                ):
+                    img = child[0]
+                    div = etree.Element(
+                        "div", {"class": "image-container"}
+                    )
+                    div.append(img)
+                    title = img.attrib.get("alt", "")
+                    if title:
+                        raw_caption = self.md.convert(title)
+                        caption_html = raw_caption.strip()
+                        if caption_html.startswith(
+                            "<p>"
+                        ) and caption_html.endswith("</p>"):
+                            caption_html = caption_html[3:-4]
+                        caption = etree.Element("small")
+                        caption.text = ""  # should be rendered
+                        caption_html_element = etree.fromstring(
+                            f"<span>{caption_html}</span>"
+                        )
+                        caption.append(caption_html_element)
+                        div.append(caption)
+                    parent[idx] = div
+
+
+class ZonaLinkTreeprocessor(Treeprocessor):
     def __init__(
         self,
         config: ZonaConfig | None,
@@ -123,12 +162,17 @@ def md_to_html(
                 "md_to_html() missing source and ctx when resolve_links is true"
             )
         md.treeprocessors.register(
-            item=ZonaLinkProcessor(
+            item=ZonaLinkTreeprocessor(
                 config, resolve_links, source, layout, item_map
             ),
             name="zona_links",
             priority=15,
         )
+    md.treeprocessors.register(
+        item=ZonaImageTreeprocessor(md),
+        name="zona_images",
+        priority=17,
+    )
     return md.convert(content)