lib/rehype-unwrap-images.tsx

@@ -1,26 +1,61 @@
 import type { Root, Element, ElementContent } from 'hast'
 import type { Plugin } from 'unified'
-import { visit } from 'unist-util-visit'
+import { SKIP, visit } from 'unist-util-visit'
 
-function isWhitespaceText(node: ElementContent) {
-  return node.type === 'text' && node.value.trim() === ''
+function isBlank(node: ElementContent) {
+  return (
+    (node.type === 'text' && node.value.trim() === '') ||
+    (node.type === 'element' && node.tagName === 'br')
+  )
+}
+
+function isImage(node: ElementContent): node is Element {
+  return node.type === 'element' && node.tagName === 'img'
 }
 
 /**
- * A rehype plugin that unwraps <img> elements from <p> tags if the <p> contains only the <img> and optional whitespace.
- * For example, a paragraph like "<p><img src='image.png' alt='An image'></p>" will be transformed into "<img src='image.png' alt='An image'>".
+ * A rehype plugin that lifts <img> elements out of <p> tags, since images render as <figure>,
+ * which can't be nested inside a paragraph.
+ *
+ * A paragraph containing only an image is replaced by the image. A paragraph mixing text and
+ * images (e.g. "**Screenshots:**" followed by images on the next lines) is split: each image
+ * becomes a standalone node and the surrounding inline content stays in its own <p>.
  */
 export const rehypeUnwrapImages: Plugin<[], Root> = () => (tree: Root) => {
   visit(tree, 'element', (node: Element, index, parent) => {
     if (node.tagName !== 'p' || !parent || typeof index !== 'number') return
+    if (!node.children.some(isImage)) return
 
-    const meaningfulChildren = node.children.filter((c) => !isWhitespaceText(c))
-    if (
-      meaningfulChildren.length === 1 &&
-      meaningfulChildren[0].type === 'element' &&
-      meaningfulChildren[0].tagName === 'img'
-    ) {
-      parent.children[index] = meaningfulChildren[0]
+    const replacement: ElementContent[] = []
+    let run: ElementContent[] = []
+
+    const flushRun = () => {
+      // Drop whitespace and line breaks left over around the lifted images
+      if (run.some((child) => !isBlank(child))) {
+        replacement.push({ ...node, children: trimBlankEdges(run) })
+      }
+      run = []
     }
+
+    for (const child of node.children) {
+      if (isImage(child)) {
+        flushRun()
+        replacement.push(child)
+      } else {
+        run.push(child)
+      }
+    }
+    flushRun()
+
+    parent.children.splice(index, 1, ...replacement)
+    return [SKIP, index + replacement.length]
   })
 }
+
+function trimBlankEdges(nodes: ElementContent[]): ElementContent[] {
+  let start = 0
+  let end = nodes.length
+  while (start < end && isBlank(nodes[start])) start++
+  while (end > start && isBlank(nodes[end - 1])) end--
+  return nodes.slice(start, end)
+}