From 6ac57d782293ece689ee7b93fc418c86bbef0ac0 Mon Sep 17 00:00:00 2001
From: InSyncWithFoo <insyncwithfoo@gmail.com>
Date: Mon, 3 Mar 2025 13:19:59 +0000
Subject: [PATCH] Allow folding script metadata block into `# /// script`

---
 .../ScriptMetadataBlockEnterHandler.kt        |  2 +-
 .../ScriptMetadataFoldingBuilder.kt           | 50 +++++++++++++++++++
 .../ScriptMetadataTemplateContext.kt          |  9 ++++
 src/main/resources/META-INF/others.xml        |  7 +++
 4 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataFoldingBuilder.kt

diff --git a/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataBlockEnterHandler.kt b/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataBlockEnterHandler.kt
index b691d30..80d180a 100644
--- a/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataBlockEnterHandler.kt
+++ b/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataBlockEnterHandler.kt
@@ -22,7 +22,7 @@ private val possiblyEmptyScriptBlock = """(?mx)
 
 
 /**
- * Insert `# ` when the user presses Enter
+ * Insert a leading `#` when the user presses Enter
  * in the middle of a script metadata block.
  * 
  * Before:
diff --git a/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataFoldingBuilder.kt b/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataFoldingBuilder.kt
new file mode 100644
index 0000000..7ea49d2
--- /dev/null
+++ b/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataFoldingBuilder.kt
@@ -0,0 +1,50 @@
+package insyncwithfoo.ryecharm.others.scriptmetadata
+
+import com.intellij.lang.ASTNode
+import com.intellij.lang.folding.CustomFoldingBuilder
+import com.intellij.lang.folding.FoldingDescriptor
+import com.intellij.openapi.editor.Document
+import com.intellij.openapi.project.DumbAware
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.startOffset
+import com.jetbrains.python.psi.PyFile
+
+
+/**
+ * Fold a script metadata block into `# /// script`.
+ */
+internal class ScriptMetadataFoldingBuilder : CustomFoldingBuilder(), DumbAware {
+    
+    override fun getLanguagePlaceholderText(node: ASTNode, range: TextRange) =
+        "# /// script"
+    
+    override fun isRegionCollapsedByDefault(node: ASTNode) = false
+    
+    override fun buildLanguageFoldRegions(
+        descriptors: MutableList<FoldingDescriptor>,
+        root: PsiElement,
+        document: Document,
+        quick: Boolean
+    ) {
+        if (root !is PyFile) {
+            return
+        }
+        
+        val block = scriptBlock.find(document.charsSequence) ?: return
+        val blockRange = TextRange(block.range.first, block.range.last + 1)
+        
+        for (element in root.children) {
+            if (element !is PsiComment) {
+                continue
+            }
+            
+            if (element.startOffset == blockRange.startOffset) {
+                descriptors += FoldingDescriptor(element, blockRange)
+                break
+            }
+        }
+    }
+    
+}
diff --git a/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataTemplateContext.kt b/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataTemplateContext.kt
index 8310dc2..2b5f544 100644
--- a/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataTemplateContext.kt
+++ b/src/main/kotlin/insyncwithfoo/ryecharm/others/scriptmetadata/ScriptMetadataTemplateContext.kt
@@ -7,6 +7,15 @@ import com.jetbrains.python.psi.LanguageLevel
 import insyncwithfoo.ryecharm.message
 
 
+/**
+ * Determine whether the `script` live template should be suggested.
+ * 
+ * The cursor is considered to be "in context" when
+ * all of the following are true:
+ * 
+ * * The first character of the current line is not a space.
+ * * The file has yet to have a script metadata block.
+ */
 internal class ScriptMetadataTemplateContext :
     TemplateContextType(message("templates.scriptMetadata.presentableName")) {
     
diff --git a/src/main/resources/META-INF/others.xml b/src/main/resources/META-INF/others.xml
index f5746c7..24bbfd7 100644
--- a/src/main/resources/META-INF/others.xml
+++ b/src/main/resources/META-INF/others.xml
@@ -34,6 +34,13 @@
 			implementation="insyncwithfoo.ryecharm.others.scriptmetadata.ScriptMetadataTemplateContext"
 		/>
 		
+		<lang.foldingBuilder
+			id="insyncwithfoo.ryecharm.others.scriptmetadata.ScriptMetadataFoldingBuilder"
+			order="first"
+			language="Python"
+			implementationClass="insyncwithfoo.ryecharm.others.scriptmetadata.ScriptMetadataFoldingBuilder"
+		/>
+		
 		<languageInjectionContributor
 			id="insyncwithfoo.ryecharm.others.requirementsinjection.RequirementsInjector"
 			language="TOML"