Skip to content

Commit 0a274bb

Browse files
PetrYakopyakovenko
PetrYako
and
pyakovenko
authored
feat(toolwindow): save state of the branches across files (#379)
Co-authored-by: pyakovenko <peotr.yakovenko@gmail.com>
1 parent de244c0 commit 0a274bb

File tree

3 files changed

+172
-5
lines changed

3 files changed

+172
-5
lines changed

‎src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import com.intellij.openapi.fileEditor.FileEditorManagerEvent
99
import com.intellij.openapi.fileEditor.FileEditorManagerListener
1010
import com.intellij.openapi.project.Project
1111
import com.intellij.openapi.ui.SimpleToolWindowPanel
12+
import com.intellij.openapi.vfs.VirtualFile
1213
import com.intellij.openapi.vfs.VirtualFileManager
1314
import com.intellij.openapi.vfs.newvfs.BulkFileListener
1415
import com.intellij.openapi.vfs.newvfs.events.VFileEvent
@@ -104,6 +105,11 @@ class TestExplorerWindow(private val project: Project) : SimpleToolWindowPanel(t
104105
project.messageBus.connect().subscribe(
105106
FileEditorManagerListener.FILE_EDITOR_MANAGER,
106107
object : FileEditorManagerListener {
108+
override fun fileClosed(source: FileEditorManager, file: VirtualFile) {
109+
// when a file is closed, reset the one-time expanded state so reopening expands all again
110+
tree.markFileClosed(file)
111+
}
112+
107113
override fun selectionChanged(event: FileEditorManagerEvent) {
108114
val file = fileEditorManager.selectedEditor?.file
109115
if (file != null) {

‎src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt‎

Lines changed: 65 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,18 @@ package io.kotest.plugin.intellij.toolwindow
33
import com.intellij.ide.util.treeView.NodeRenderer
44
import com.intellij.ide.util.treeView.PresentableNodeDescriptor
55
import com.intellij.openapi.project.Project
6+
import com.intellij.openapi.vfs.VirtualFile
67
import com.intellij.ui.TreeUIHelper
78
import javax.swing.tree.DefaultMutableTreeNode
89
import javax.swing.tree.TreeModel
910
import javax.swing.tree.TreeSelectionModel
1011

12+
private data class FileTreeState(
13+
val allKeys: Set<String>,
14+
val expandedKeys: Set<String>,
15+
var initiallyExpanded: Boolean,
16+
)
17+
1118
class TestFileTree(
1219
project: Project,
1320
) : com.intellij.ui.treeStructure.Tree(),
@@ -17,6 +24,8 @@ class TestFileTree(
1724
private val kotestTestExplorerService: KotestTestExplorerService =
1825
project.getService(KotestTestExplorerService::class.java)
1926
private var initialized = false
27+
private var lastFileKey: String? = null
28+
private val stateByFileKey = mutableMapOf<String, FileTreeState>()
2029

2130
init {
2231
selectionModel.selectionMode = TreeSelectionModel.SINGLE_TREE_SELECTION
@@ -40,13 +49,64 @@ class TestFileTree(
4049
super.setModel(treeModel)
4150
return
4251
}
43-
val expanded = isExpanded(0)
52+
val newFileKey = currentFileKey()
53+
54+
// If switching away from a file, save its state first
55+
if (lastFileKey != null && lastFileKey != newFileKey) {
56+
val prevAll = collectAllPathKeys()
57+
val prevExpanded = collectExpandedPathKeys()
58+
val prevInit = stateByFileKey[lastFileKey!!]?.initiallyExpanded ?: false
59+
stateByFileKey[lastFileKey!!] = FileTreeState(prevAll, prevExpanded, prevInit)
60+
}
61+
62+
val sameFile = newFileKey == lastFileKey
63+
val prevStateForNew = if (newFileKey != null) stateByFileKey[newFileKey] else null
64+
val firstOpenForFile = newFileKey != null && prevStateForNew == null
65+
66+
// Baselines (use live tree for same file; fallback to stored state when switching)
67+
val prevAllKeysForThisFile: Set<String> = when {
68+
firstOpenForFile -> emptySet()
69+
sameFile -> collectAllPathKeys()
70+
newFileKey != null -> prevStateForNew?.allKeys ?: emptySet()
71+
else -> emptySet()
72+
}
73+
val expandedKeysToRestore: Set<String> = when {
74+
firstOpenForFile -> emptySet()
75+
sameFile -> collectExpandedPathKeys()
76+
newFileKey != null -> prevStateForNew?.expandedKeys ?: emptySet()
77+
else -> emptySet()
78+
}
79+
4480
super.setModel(treeModel)
45-
expandAllNodes()
46-
setModuleGroupNodeExpandedState(expanded)
81+
82+
// Compute added nodes relative to the previous snapshot of this file (if any)
83+
val newAllKeys = collectAllPathKeys()
84+
if (!firstOpenForFile) {
85+
val addedKeys = newAllKeys - prevAllKeysForThisFile
86+
if (addedKeys.isNotEmpty()) expandAncestorPrefixesFor(addedKeys)
87+
}
88+
89+
if (firstOpenForFile) {
90+
// First time this file is shown in the tool window: expand everything
91+
expandAllNodes()
92+
stateByFileKey[newFileKey] = FileTreeState(newAllKeys, collectExpandedPathKeys(), initiallyExpanded = true)
93+
} else {
94+
// Restore previous expansion state for this file
95+
if (expandedKeysToRestore.isNotEmpty()) expandPathsByKeys(expandedKeysToRestore)
96+
if (newFileKey != null) {
97+
val init = prevStateForNew?.initiallyExpanded ?: true
98+
stateByFileKey[newFileKey] = FileTreeState(newAllKeys, collectExpandedPathKeys(), init)
99+
}
100+
}
101+
102+
lastFileKey = newFileKey
47103
}
48104

49-
private fun setModuleGroupNodeExpandedState(expanded: Boolean) {
50-
if (expanded) expandRow(0) else collapseRow(0)
105+
fun markFileClosed(file: VirtualFile) {
106+
stateByFileKey.remove(file.path)
107+
if (lastFileKey == file.path) lastFileKey = null
51108
}
109+
110+
private fun currentFileKey(): String? = kotestTestExplorerService.currentFile?.path
111+
52112
}

‎src/main/kotlin/io/kotest/plugin/intellij/toolwindow/treeutils.kt‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,3 +34,104 @@ fun TreePath.nodeDescriptor(): PresentableNodeDescriptor<*>? {
3434
else -> null
3535
}
3636
}
37+
38+
// --- Expansion state helpers ---
39+
40+
/**
41+
* Builds a logical key for a node to preserve expansion state across model rebuilds.
42+
* Only nodes that can have children are keyed (root, modules, module, tags, file, spec, container test).
43+
*/
44+
private fun DefaultMutableTreeNode.expansionKeyOrNull(): String? {
45+
return when (val descriptor = userObject) {
46+
is KotestRootNodeDescriptor -> "root"
47+
is ModulesNodeDescriptor -> "modules"
48+
is ModuleNodeDescriptor -> "module:${descriptor.module.name}"
49+
is TagsNodeDescriptor -> "tags"
50+
is TestFileNodeDescriptor -> "file"
51+
is SpecNodeDescriptor -> "spec:${descriptor.fqn.asString()}"
52+
is TestNodeDescriptor -> {
53+
// Only container tests can have children; still safe to key all tests
54+
"test:${descriptor.test.test.descriptorPath()}"
55+
}
56+
else -> null
57+
}
58+
}
59+
60+
/**
61+
* Returns a set of expansion path keys for the currently expanded nodes.
62+
*/
63+
fun JTree.collectExpandedPathKeys(): Set<String> {
64+
val root = model.root as? DefaultMutableTreeNode ?: return emptySet()
65+
val expanded = mutableSetOf<String>()
66+
// Enumerate all expanded descendants starting from root
67+
val enumeration = getExpandedDescendants(TreePath(root.path)) ?: return emptySet()
68+
while (enumeration.hasMoreElements()) {
69+
val path = enumeration.nextElement()
70+
val key = pathToExpansionKey(path)
71+
if (key != null) expanded.add(key)
72+
}
73+
return expanded
74+
}
75+
76+
/**
77+
* Expands nodes in the current model whose logical expansion keys appear in [keys].
78+
*/
79+
fun JTree.expandPathsByKeys(keys: Set<String>) {
80+
val root = model.root as? DefaultMutableTreeNode ?: return
81+
fun recurse(node: DefaultMutableTreeNode, prefix: String?) {
82+
val key = node.expansionKeyOrNull()
83+
val pathKey = if (key == null) prefix else listOfNotNull(prefix, key).joinToString("/")
84+
if (pathKey != null && keys.contains(pathKey)) {
85+
expandPath(TreePath(node.path))
86+
}
87+
val children = node.children()
88+
while (children.hasMoreElements()) {
89+
val child = children.nextElement() as DefaultMutableTreeNode
90+
recurse(child, pathKey)
91+
}
92+
}
93+
recurse(root, null)
94+
}
95+
96+
private fun pathToExpansionKey(path: TreePath): String? {
97+
val parts = path.path
98+
.mapNotNull { it as? DefaultMutableTreeNode }
99+
.mapNotNull { it.expansionKeyOrNull() }
100+
return if (parts.isEmpty()) null else parts.joinToString("/")
101+
}
102+
103+
/**
104+
* Returns a set of path keys for all nodes in the current model.
105+
*/
106+
fun JTree.collectAllPathKeys(): Set<String> {
107+
val root = model.root as? DefaultMutableTreeNode ?: return emptySet()
108+
val keys = mutableSetOf<String>()
109+
fun recurse(node: DefaultMutableTreeNode) {
110+
val path = TreePath(node.path)
111+
val key = pathToExpansionKey(path)
112+
if (key != null) keys.add(key)
113+
val children = node.children()
114+
while (children.hasMoreElements()) {
115+
recurse(children.nextElement() as DefaultMutableTreeNode)
116+
}
117+
}
118+
recurse(root)
119+
return keys
120+
}
121+
122+
/**
123+
* Expands all ancestor prefixes for the given full path keys.
124+
*/
125+
fun JTree.expandAncestorPrefixesFor(keys: Set<String>) {
126+
if (keys.isEmpty()) return
127+
val toExpand = mutableSetOf<String>()
128+
keys.forEach { key ->
129+
val parts = key.split('/')
130+
val acc = mutableListOf<String>()
131+
parts.forEach { part ->
132+
acc.add(part)
133+
toExpand.add(acc.joinToString("/"))
134+
}
135+
}
136+
expandPathsByKeys(toExpand)
137+
}

0 commit comments

Comments
 (0)