From 4df195470851b1a881b6ac87e7031d156f68dff5 Mon Sep 17 00:00:00 2001
From: Adam Turner <9087854+aa-turner@users.noreply.github.com>
Date: Sun, 13 Nov 2022 19:04:03 +0000
Subject: [PATCH] Refactor to getHideableCopyButtonElements

---
 python_docs_theme/static/copybutton.js | 52 +++++++++++++++++---------
 1 file changed, 34 insertions(+), 18 deletions(-)

diff --git a/python_docs_theme/static/copybutton.js b/python_docs_theme/static/copybutton.js
index 81e68c5..5068132 100644
--- a/python_docs_theme/static/copybutton.js
+++ b/python_docs_theme/static/copybutton.js
@@ -1,3 +1,31 @@
+function* getHideableCopyButtonElements(rootElement) {
+    // yield all elements with the "go" (Generic.Output),
+    // "gp" (Generic.Prompt), or "gt" (Generic.Traceback) CSS class
+    for (const el of rootElement.querySelectorAll('.go, .gp, .gt')) {
+        yield el
+    }
+    // tracebacks (.gt) contain bare text elements that need to be
+    // wrapped in a span to hide or show the element
+    for (let el of rootElement.querySelectorAll('.gt')) {
+        while ((el = el.nextSibling) && el.nodeType !== Node.DOCUMENT_NODE) {
+            // stop wrapping text nodes when we hit the next output or
+            // prompt element
+            if (el.nodeType === Node.ELEMENT_NODE && el.matches(".gp, .go")) {
+                break
+            }
+            // if the node is a text node with content, wrap it in a
+            // span element so that we can control visibility
+            if (el.nodeType === Node.TEXT_NODE && el.textContent.trim()) {
+                const wrapper = document.createElement('span')
+                el.after(wrapper)
+                wrapper.appendChild(el)
+                el = wrapper
+            }
+            yield el
+        }
+    }
+}
+
 
 const loadCopyButton = () => {
     /* Add a [>>>] button in the top-right corner of code samples to hide
@@ -18,28 +46,16 @@ const loadCopyButton = () => {
         const codeEl = buttonEl.nextElementSibling
         if (buttonEl.dataset.hidden === 'false') {
             // hide the code output
-            codeEl.querySelectorAll('.go, .gp, .gt').forEach(el => el.hidden = true)
-            // tracebacks (.gt) contain bare text elements that need to be
-            // wrapped in a span to hide or show the element
-            codeEl.querySelectorAll('.gt').forEach(el => {
-                while ((el = el.nextSibling) && el.nodeType !== Node.DOCUMENT_NODE) {
-                    if (el.nodeType === Node.ELEMENT_NODE && el.matches(".gp, .go")) {
-                        break
-                    }
-                    if (el.nodeType === Node.TEXT_NODE && el.textContent.trim()) {
-                        const wrapper = document.createElement('span')
-                        el.after(wrapper)
-                        wrapper.appendChild(el)
-                        el = wrapper
-                    }
-                    el.hidden = true
-                }
-            })
+            for (const el of getHideableCopyButtonElements(codeEl)) {
+                el.hidden = true
+            }
             buttonEl.title = show_text
             buttonEl.dataset.hidden = "true"
         } else {
             // show the code output
-            codeEl.childNodes.forEach(el => el.hidden = false)
+            for (const el of getHideableCopyButtonElements(codeEl)) {
+                el.hidden = false
+            }
             buttonEl.title = hide_text
             buttonEl.dataset.hidden = "false"
         }