fix: fold TSX object literal attributes

Fixes #14

Author: Thomas Scharke

src/main/kotlin/dev/zbinski/htmlattributefolder/AttributeFolder.kt

@@ -20,13 +20,53 @@ class AttributeFolder: FoldingBuilderEx(), DumbAware {
         for (item in getAttributes(Array(1) { root })) {
             var end: Int
             var start: Int
+
             if (settings.foldingMethod == 1) {
                 end = item.attribute.textRange.endOffset
                 start = item.attribute.textRange.startOffset
             } else {
-                val len = item.attributeName.length + settings.attributeSeparator.length + settings.attributeWrapper.length
-                start = item.attribute.textRange.startOffset + len
-                end = item.attribute.textRange.endOffset - settings.attributeWrapper.length
+                val text = item.attribute.text
+                val base = item.attribute.textRange.startOffset
+
+                // Find '=' then the first non-space char after it
+                val eq = text.indexOf(settings.attributeSeparator)
+                if (eq < 0) continue
+                var i = eq + 1
+                while (i < text.length && text[i].isWhitespace()) i++
+                if (i >= text.length) continue
+
+                when (text[i]) {
+                    // name="value" or name='value'
+                    '"', '\'' -> {
+                        val quote = text[i]
+                        val open = i
+                        val close = text.indexOf(quote, startIndex = open + 1)
+                        if (close < 0) continue
+
+                        start = base + open + 1
+                        end = base + close
+                    }
+                    // name={...} or name={{...}}
+                    '{' -> {
+                        val outerOpen = i
+                        val outerClose = findMatchingBrace(text, outerOpen) ?: continue
+
+                        val isDouble = outerOpen + 1 < text.length && text[outerOpen + 1] == '{'
+                        if (isDouble) {
+                            val innerOpen = outerOpen + 1
+                            val innerClose = findMatchingBrace(text, innerOpen) ?: continue
+
+                            // Fold inside INNER braces INCLUDING whitespace so output becomes: style={{__PLACEHOLDER__}}
+                            start = base + innerOpen + 1
+                            end = base + innerClose
+                        } else {
+                            // Fold inside single braces INCLUDING whitespace: name={__PLACEHOLDER__}
+                            start = base + outerOpen + 1
+                            end = base + outerClose
+                        }
+                    }
+                    else -> continue
+                }
             }
 
             if (end > start) {
@@ -57,24 +97,46 @@ class AttributeFolder: FoldingBuilderEx(), DumbAware {
         return settings.collapseByDefault
     }
 
+    private fun findMatchingBrace(text: String, openIndex: Int): Int? {
+        if (openIndex !in text.indices || text[openIndex] != '{') return null
+        var depth = 0
+        var i = openIndex
+        while (i < text.length) {
+            when (text[i]) {
+                '{' -> depth++
+                '}' -> {
+                    depth--
+                    if (depth == 0) return i
+                }
+            }
+            i++
+        }
+        return null
+    }
+
     private fun getAttributes(
         elements: Array<PsiElement>,
         attributes: ArrayList<String> = settings.attributes
     ): Sequence<Attribute> = sequence {
         for (child in elements) {
+            val t = child.text
             for (attributeName in attributes) {
-                val attributeBeginning = attributeName + settings.attributeSeparator + settings.attributeWrapper
-                if (child.text.startsWith(attributeBeginning)) {
+                val startsLikeAttribute =
+                    t.startsWith("$attributeName=\"") ||
+                            t.startsWith("$attributeName='") ||
+                            t.startsWith("$attributeName={")
+
+                if (startsLikeAttribute) {
                     yield(object : Attribute {
                         override val attribute = child
                         override val attributeName = attributeName
                     })
                 }
+            }
 
-                val items = getAttributes(child.children, arrayListOf(attributeName)).iterator()
-                while (items.hasNext()) {
-                    yield(items.next())
-                }
+            val items = getAttributes(child.children, attributes).iterator()
+            while (items.hasNext()) {
+                yield(items.next())
             }
         }
     }

src/test/kotlin/dev.zbinski.htmlattributefolder/AttributeFolderTSXTest.kt

@@ -2,38 +2,38 @@ package dev.zbinski.htmlattributefolder
 
 private const val TSX_SNIPPET = """
     export const Component = () => (
-      <div class="a b" className="c d" style={{ backgroundColor: "red" }} data-foo="foo-data" />
+      <div class="a b" className="c d" style={{ backgroundColor: "red", nested: { "foo": "bar" } }} data-foo="foo-data" />
     )
 """
 
 class AttributeFolderTSXTest : BaseAttributeFolderTest() {
     fun testFoldingTSXAttributesCollapsed() {
         assertContains(document.text, """class="a b"""")
         assertContains(document.text, """className="c d"""")
-        assertContains(document.text, """style={{ backgroundColor: "red" }}""")
+        assertContains(document.text, """style={{ backgroundColor: "red", nested: { "foo": "bar" } }}""")
         assertContains(document.text, """data-foo="foo-data"""")
 
-        configureAttributeFolder(collapseByDefault = true)
+        configureAttributeFolder(collapseByDefault = true, arrayListOf("class", "className", "style"))
 
         val visualText = applyPluginFoldingAndRender()
         assertContains(visualText, """class="__PLACEHOLDER__"""")
         assertContains(visualText, """className="__PLACEHOLDER__"""")
-        assertContains(visualText, """style={{ backgroundColor: "red" }}""")
+        assertContains(visualText, """style={{__PLACEHOLDER__}}""")
         assertContains(visualText, """data-foo="foo-data"""")
     }
 
     fun testFoldingTSXAttributesUncollapsed() {
         assertContains(document.text, """class="a b"""")
         assertContains(document.text, """className="c d"""")
-        assertContains(document.text, """style={{ backgroundColor: "red" }}""")
+        assertContains(document.text, """style={{ backgroundColor: "red", nested: { "foo": "bar" } }}""")
         assertContains(document.text, """data-foo="foo-data"""")
 
-        configureAttributeFolder(collapseByDefault = false)
+        configureAttributeFolder(collapseByDefault = false, arrayListOf("class", "className", "style"))
 
         val visualText = applyPluginFoldingAndRender()
         assertContains(visualText, """class="a b"""")
         assertContains(visualText, """className="c d"""")
-        assertContains(visualText, """style={{ backgroundColor: "red" }}""")
+        assertContains(visualText, """style={{ backgroundColor: "red", nested: { "foo": "bar" } }}""")
         assertContains(visualText, """data-foo="foo-data"""")
     }