Repository: stateful/vscode-runme
Pull request #145 (2023-01-25 01:01:44)

Implement collapse/expand all buttons to explorer
implements #130

christian-bromann: @sourishkrout @degrammer this is ready for review now, here is a demo:

![tree](https://user-images.githubusercontent.com/731337/214716839-bb67ebe4-ba4b-4a3c-906d-69118ab09214.gif)

It took a bit longer as VS Code seem to show some weird behavior with `TreeView`s, read more on this here: https://github.com/microsoft/vscode/issues/172479

Let me know if anything is unclear.

Patch:
diff --git a/package.json b/package.json
index 4a41a890..9b79e50b 100644
--- a/package.json
+++ b/package.json
@@ -101,6 +101,18 @@
         "command": "runme.stopBackgroundTask",
         "category": "Runme",
         "title": "Dispose terminal session to stop background task."
+      },
+      {
+        "command": "runme.collapseTreeView",
+        "category": "Runme",
+        "title": "Collapse all items",
+        "icon": "$(debug-collapse-all)"
+      },
+      {
+        "command": "runme.expandTreeView",
+        "category": "Runme",
+        "title": "Expand all items",
+        "icon": "$(search-expand-results)"
       }
     ],
     "menus": {
@@ -120,6 +132,14 @@
         {
           "command": "runme.stopBackgroundTask",
           "when": "false"
+        },
+        {
+          "command": "runme.collapseTreeView",
+          "when": "false"
+        },
+        {
+          "command": "runme.expandTreeView",
+          "when": "false"
         }
       ],
       "editor/title": [
@@ -139,6 +159,18 @@
           "command": "runme.new",
           "group": "notebook"
         }
+      ],
+      "view/title": [
+        {
+          "command": "runme.collapseTreeView",
+          "when": "view == runme.launcher && runme.launcher.isExpanded",
+          "group": "navigation"
+        },
+        {
+          "command": "runme.expandTreeView",
+          "when": "view == runme.launcher && !runme.launcher.isExpanded",
+          "group": "navigation"
+        }
       ]
     },
     "notebooks": [
diff --git a/src/extension/extension.ts b/src/extension/extension.ts
index abe2fc30..801261b3 100644
--- a/src/extension/extension.ts
+++ b/src/extension/extension.ts
@@ -27,6 +27,8 @@ export class RunmeExtension {
     const kernel = new Kernel(context)
     const grpcSerializer = kernel.hasExperimentEnabled('grpcSerializer')
     const serializer = grpcSerializer ? new GrpcSerializer(context) : new WasmSerializer(context)
+    const treeViewer = new RunmeLauncherProvider(getDefaultWorkspace())
+
     context.subscriptions.push(
       kernel,
       workspace.registerNotebookSerializer(Kernel.type, serializer, {
@@ -53,7 +55,13 @@ export class RunmeExtension {
       RunmeExtension.registerCommand('runme.new', createNewRunmeNotebook),
       RunmeExtension.registerCommand('runme.openRunmeFile', RunmeLauncherProvider.openFile),
       tasks.registerTaskProvider(RunmeTaskProvider.id, new RunmeTaskProvider(context)),
-      window.registerTreeDataProvider('runme.launcher', new RunmeLauncherProvider(getDefaultWorkspace()))
+
+      /**
+       * tree viewer items
+       */
+      window.registerTreeDataProvider('runme.launcher', treeViewer),
+      commands.registerCommand('runme.collapseTreeView', treeViewer.collapseAll.bind(treeViewer)),
+      commands.registerCommand('runme.expandTreeView', treeViewer.expandAll.bind(treeViewer))
     )
 
     /**
diff --git a/src/extension/provider/launcher.ts b/src/extension/provider/launcher.ts
index 4e8a76cf..a3a58993 100644
--- a/src/extension/provider/launcher.ts
+++ b/src/extension/provider/launcher.ts
@@ -25,78 +25,112 @@ interface IRunmeFileProps {
   contextValue: string
 }
 
+interface TreeFile {
+  files: string[]
+  folderPath: string
+}
+
+/**
+ * used to force VS Code update the tree view when user expands/collapses all
+ * see https://github.com/microsoft/vscode/issues/172479
+ */
+let i = 0
+
 export class RunmeFile extends TreeItem {
-  constructor(
-    public readonly label: string,
-    { collapsibleState, tooltip, onSelectedCommand, lightIcon, darkIcon, contextValue }: IRunmeFileProps
-  ) {
-    super(label, collapsibleState)
+  constructor(public label: string, options: IRunmeFileProps) {
+    super(label, options.collapsibleState)
     const assetsPath = join(__filename, '..', '..', 'assets')
-    this.tooltip = tooltip
-    this.label = label
-    this.command = onSelectedCommand
-    this.contextValue = contextValue
+
+    this.tooltip = options.tooltip
+    this.command = options.onSelectedCommand
+    this.contextValue = options.contextValue
     this.iconPath = {
-      light: join(assetsPath, lightIcon),
-      dark: join(assetsPath, darkIcon),
+      light: join(assetsPath, options.lightIcon),
+      dark: join(assetsPath, options.darkIcon),
     }
   }
 }
 
 export class RunmeLauncherProvider implements TreeDataProvider<RunmeFile> {
-  private filesTree: Map<string, any>
-  constructor(private workspaceRoot?: string | undefined) {
-    this.filesTree = new Map()
-  }
+  private filesTree: Map<string, TreeFile> = new Map()
+  private defaultItemState: TreeItemCollapsibleState = TreeItemCollapsibleState.Collapsed
 
-  private _onDidChangeTreeData: EventEmitter<RunmeFile | undefined | void> = new EventEmitter<
-    RunmeFile | undefined | void
-  >()
+  constructor(private workspaceRoot?: string | undefined) {}
 
-  readonly onDidChangeTreeData: Event<RunmeFile | undefined | void> = this._onDidChangeTreeData.event
-  getTreeItem(element: RunmeFile): TreeItem | Thenable<TreeItem> {
+  private _onDidChangeTreeData: EventEmitter<RunmeFile | undefined> = new EventEmitter<RunmeFile | undefined>()
+  readonly onDidChangeTreeData: Event<RunmeFile | undefined> = this._onDidChangeTreeData.event
+
+  getTreeItem(element: RunmeFile): TreeItem {
     return element
   }
-  getChildren(element?: RunmeFile | undefined): Thenable<RunmeFile[]> {
+
+  async getChildren(element?: RunmeFile | undefined): Promise<RunmeFile[]> {
     if (!this.workspaceRoot) {
       return Promise.resolve([])
     }
 
-    if (!element) {
-      return new Promise(async (resolve: (value: RunmeFile[]) => void) => {
-        await this.getRunmeFilesFromWorkspace(resolve)
-      })
+    /**
+     * scan for files first time we render the tree
+     */
+    if (this.filesTree.size === 0) {
+      await this._scanWorkspace()
     }
 
-    const { files, folderPath } = this.filesTree.get(element.label)
-    const folderMarkdownItems: RunmeFile[] = []
-
-    for (const file of files) {
-      folderMarkdownItems.push(
-        new RunmeFile(file, {
-          collapsibleState: TreeItemCollapsibleState.None,
-          tooltip: 'Click to open runme file',
-          lightIcon: 'icon.gif',
-          darkIcon: 'icon.gif',
-          contextValue: 'markdown-file',
-          onSelectedCommand: {
-            arguments: [{ file, folderPath }],
-            command: 'runme.openRunmeFile',
-            title: file,
-          },
-        })
-      )
+    if (!element) {
+      /**
+       * render folders
+       */
+      return [...this.filesTree.keys()].map((folder) => new RunmeFile(folder, {
+        collapsibleState: this.defaultItemState,
+        tooltip: 'Click to open runme files from folder',
+        lightIcon: 'folder.svg',
+        darkIcon: 'folder.svg',
+        contextValue: 'folder',
+      })).sort((a: RunmeFile, b: RunmeFile) => a.label.length > b.label.length ? 1 : -1)
     }
 
-    return Promise.resolve(folderMarkdownItems)
+    const { files, folderPath } = this.filesTree.get(element.label as string) || { files: [] }
+    return files.map((file) => new RunmeFile(file, {
+      collapsibleState: TreeItemCollapsibleState.None,
+      tooltip: 'Click to open runme file',
+      lightIcon: 'icon.gif',
+      darkIcon: 'icon.gif',
+      contextValue: 'markdown-file',
+      onSelectedCommand: {
+        arguments: [{ file, folderPath }],
+        command: 'runme.openRunmeFile',
+        title: file,
+      }
+    }))
+  }
+
+  getParent () {
+    return null
   }
 
   refresh(): void {
-    this._onDidChangeTreeData.fire()
+    this.filesTree = new Map()
+    this._onDidChangeTreeData.fire(undefined)
+  }
+
+  public static async openFile({ file, folderPath }: { file: string, folderPath: string }) {
+    const doc = Uri.file(`${folderPath}/${file}`)
+    await commands.executeCommand('vscode.openWith', doc, Kernel.type)
+  }
+
+  async collapseAll () {
+    this.defaultItemState = TreeItemCollapsibleState.Collapsed
+    await commands.executeCommand('setContext', 'runme.launcher.isExpanded', false)
+    this.refresh()
+  }
+
+  async expandAll () {
+    this.defaultItemState = TreeItemCollapsibleState.Expanded
+    await commands.executeCommand('setContext', 'runme.launcher.isExpanded', true)
+    this.refresh()
   }
 
-  async getRunmeFilesFromWorkspace(onComplete: (value: RunmeFile[]) => void): Promise<void> {
-    const runmeFileCollection: RunmeFile[] = []
+  private async _scanWorkspace () {
     let excludePatterns
 
     if (this.workspaceRoot) {
@@ -112,34 +146,21 @@ export class RunmeLauncherProvider implements TreeDataProvider<RunmeFile> {
     const files = await workspace.findFiles('**/*.md', `{${excludePatterns}}`)
     const rootFolder = basename(this.workspaceRoot || '')
 
+    /**
+     * we need to tweak the folder name to force VS Code re-render the tree view
+     * see https://github.com/microsoft/vscode/issues/172479
+     */
+    const nameTweaker = ++i % 2 ? ' ' : ''
     for (const { path } of files) {
       const info = basename(path)
       const folderPath = dirname(path)
-      const folderName = dirname(path).replace(resolve(this.workspaceRoot || '', '..'), '') || rootFolder
+      const folderName = dirname(path).replace(resolve(this.workspaceRoot || '', '..'), '') + nameTweaker || rootFolder
       if (!this.filesTree.has(folderName)) {
         this.filesTree.set(folderName, { files: [info], folderPath })
       } else {
-        const { files } = this.filesTree.get(folderName)
+        const { files } = this.filesTree.get(folderName) || { files: [] }
         this.filesTree.set(folderName, { files: [...files, info], folderPath })
       }
     }
-
-    for (const folder of this.filesTree.keys()) {
-      runmeFileCollection.push(
-        new RunmeFile(folder || rootFolder, {
-          collapsibleState: TreeItemCollapsibleState.Collapsed,
-          tooltip: 'Click to open runme files from folder',
-          lightIcon: 'folder.svg',
-          darkIcon: 'folder.svg',
-          contextValue: 'folder',
-        })
-      )
-    }
-    onComplete(runmeFileCollection.sort((a: RunmeFile,b: RunmeFile) => a.label.length > b.label.length ? 1 : -1 ))
-  }
-
-  public static async openFile({ file, folderPath }: { file: string, folderPath: string }) {
-    const doc = Uri.file(`${folderPath}/${file}`)
-    await commands.executeCommand('vscode.openWith', doc, Kernel.type)
   }
 }


Test patch:
diff --git a/tests/extension/extension.test.ts b/tests/extension/extension.test.ts
index 85965852..32cf2ac8 100644
--- a/tests/extension/extension.test.ts
+++ b/tests/extension/extension.test.ts
@@ -16,6 +16,6 @@ test('initializes all providers', async () => {
   await ext.initialize(context)
   expect(notebooks.registerNotebookCellStatusBarItemProvider).toBeCalledTimes(6)
   expect(workspace.registerNotebookSerializer).toBeCalledTimes(1)
-  expect(commands.registerCommand).toBeCalledTimes(10)
+  expect(commands.registerCommand).toBeCalledTimes(12)
   expect(window.registerTreeDataProvider).toBeCalledTimes(1)
 })
diff --git a/tests/extension/provider/launcher.test.ts b/tests/extension/provider/launcher.test.ts
index 3ed118af..08a102cd 100644
--- a/tests/extension/provider/launcher.test.ts
+++ b/tests/extension/provider/launcher.test.ts
@@ -1,10 +1,9 @@
-import { vi, describe, it, expect } from 'vitest'
-import {  commands } from 'vscode'
+import { vi, describe, it, expect, beforeEach } from 'vitest'
+import { commands } from 'vscode'
 
 import { RunmeFile, RunmeLauncherProvider } from '../../../src/extension/provider/launcher'
 import { getDefaultWorkspace } from '../../../src/extension/utils'
 
-
 vi.mock('vscode')
 vi.mock('vscode-telemetry')
 
@@ -27,6 +26,10 @@ vi.mock('../../../src/extension/utils', () => ({
   ])
 }))
 
+beforeEach(() => {
+  vi.mocked(commands.executeCommand).mockClear()
+})
+
 describe('Runme Notebooks', () => {
   it('returns an empty tree for empty workspace', async () => {
     const launchProvider = new RunmeLauncherProvider()
@@ -42,7 +45,8 @@ describe('Runme Notebooks', () => {
     })
 
     it('should return the items for the selected folder', async () => {
-      const element = new RunmeFile('runme/workspace/src', {
+      // space after "src" is important here, see name tweaker var in launcher
+      const element = new RunmeFile('runme/workspace/src ', {
         collapsibleState: 0,
         tooltip: 'Click to open runme file',
         lightIcon: 'icon.gif',
@@ -65,4 +69,20 @@ describe('Runme Notebooks', () => {
       expect(commands.executeCommand).toBeCalledWith('vscode.openWith', expect.any(String), 'runme')
     })
   })
+
+  it('has a expandAll method', async () => {
+    const launchProvider = new RunmeLauncherProvider()
+    launchProvider.refresh = vi.fn()
+    await launchProvider.collapseAll()
+    expect(commands.executeCommand).toBeCalledWith('setContext', 'runme.launcher.isExpanded', false)
+    expect(launchProvider.refresh).toBeCalledTimes(1)
+  })
+
+  it('has a expandAll method', async () => {
+    const launchProvider = new RunmeLauncherProvider()
+    launchProvider.refresh = vi.fn()
+    await launchProvider.expandAll()
+    expect(commands.executeCommand).toBeCalledWith('setContext', 'runme.launcher.isExpanded', true)
+    expect(launchProvider.refresh).toBeCalledTimes(1)
+  })
 })