From b10ee903b338bd40387e8380ae3dd03fcb79a598 Mon Sep 17 00:00:00 2001
From: Ryan Paul <rpaul@stripe.com>
Date: Tue, 9 Jan 2024 09:29:55 -0500
Subject: [PATCH] Improves the partial panel (#15)

---
 client/extension.ts | 16 +++++++++++++---
 client/package.json |  2 +-
 client/partials.ts  | 45 +++++++++++++++++++++++++++++++++------------
 package.json        |  2 +-
 server/package.json |  2 +-
 5 files changed, 49 insertions(+), 18 deletions(-)

diff --git a/client/extension.ts b/client/extension.ts
index 3bcce35..b042d0e 100644
--- a/client/extension.ts
+++ b/client/extension.ts
@@ -13,6 +13,7 @@ export default class Extension {
   private partials: PartialTreeProvider;
   private status: StatusItem;
   private preview: Preview;
+  private active?: Client;
 
   constructor(private context: VSC.ExtensionContext) {
     this.partials = new PartialTreeProvider();
@@ -86,6 +87,9 @@ export default class Extension {
     if (clientsToStart.length < 1) return;
     this.status.setClient(clientsToStart[0]);
     await Promise.allSettled(clientsToStart.map((client) => client.start()));
+
+    const uri = VSC.window.activeTextEditor?.document.uri;
+    if (this.active && uri) this.updatePartialsPane(this.active, uri);
   }
 
   async restart() {
@@ -126,7 +130,9 @@ export default class Extension {
       client?.canPreview()
     );
 
-    if (client) this.status.setClient(client);
+    if (!client) return;
+    this.status.setClient(client);
+    this.active = client;
   }
 
   async onActive(editor?: VSC.TextEditor) {
@@ -140,11 +146,15 @@ export default class Extension {
     if (!client) return;
     await client.start();
 
-    const output = await client.getDependencies(uri);
-    if (output) this.partials.update(output, client.uri);
+    this.updatePartialsPane(client, uri);
     this.updatePreview(client, uri);
   }
 
+  async updatePartialsPane(client: Client, uri: VSC.Uri) {
+    const output = await client.getDependencies(uri);
+    this.partials.update(output ?? undefined, client.uri);
+  }
+
   async onNewFileFromTemplate() {
     const validStates = [ClientState.Running, ClientState.Initializing];
     const options: TemplatePickerItem[] = [];
diff --git a/client/package.json b/client/package.json
index 51338d3..b164e6e 100644
--- a/client/package.json
+++ b/client/package.json
@@ -3,7 +3,7 @@
   "description": "Markdoc Extension",
   "author": "Ryan Paul",
   "license": "MIT",
-  "version": "0.0.11",
+  "version": "0.0.12",
   "scripts": {
     "build": "esbuild index.ts --bundle --outdir=dist --sourcemap=linked --external:vscode --platform=node --format=cjs",
     "build:server": "esbuild server.ts --bundle --outdir=dist --sourcemap=linked --external:vscode --platform=node --format=cjs"
diff --git a/client/partials.ts b/client/partials.ts
index 2ffdecc..211fd70 100644
--- a/client/partials.ts
+++ b/client/partials.ts
@@ -1,8 +1,11 @@
 import * as VSC from "vscode";
 import { DependencyInfo, PartialReference } from "../server/types";
 
+type PartialTreeRoot = { title: string; children?: PartialReference[] };
+type PartialTreeItem = PartialTreeRoot | PartialReference;
+
 export default class PartialTreeProvider
-  implements VSC.TreeDataProvider<PartialReference>
+  implements VSC.TreeDataProvider<PartialTreeItem>
 {
   private info?: DependencyInfo;
   private uri?: VSC.Uri;
@@ -10,14 +13,24 @@ export default class PartialTreeProvider
   private updateEmitter = new VSC.EventEmitter<void>();
   readonly onDidChangeTreeData = this.updateEmitter.event;
 
-  update(info: DependencyInfo, uri: VSC.Uri) {
+  update(info: DependencyInfo | undefined, uri: VSC.Uri) {
     this.info = info;
     this.uri = uri;
     this.updateEmitter.fire();
   }
 
-  getTreeItem(element: PartialReference): VSC.TreeItem {
+  getTreeItem(element: PartialTreeItem): VSC.TreeItem {
     if (!this.uri) return {};
+
+    if ("title" in element) {
+      const item = new VSC.TreeItem(element.title);
+      item.collapsibleState =
+        element?.children && element.children.length > 0
+          ? VSC.TreeItemCollapsibleState.Expanded
+          : VSC.TreeItemCollapsibleState.None;
+      return item;
+    }
+
     const uri = VSC.Uri.joinPath(this.uri, element.file);
     const item = new VSC.TreeItem(element.file);
     item.command = { command: "vscode.open", title: "Open", arguments: [uri] };
@@ -31,14 +44,22 @@ export default class PartialTreeProvider
   }
 
   getChildren(
-    element?: PartialReference
-  ): VSC.ProviderResult<PartialReference[]> {
-    return element
-      ? element.children
-      : !this.info
-      ? null
-      : this.info.dependents.length > 0
-      ? this.info.dependents
-      : this.info.dependencies;
+    element?: PartialTreeItem
+  ): VSC.ProviderResult<PartialTreeItem[]> {
+    if (element) return element.children;
+    if (!this.info) return [{ title: "Loading..." }];
+
+    const rootElements = [];
+    const { dependencies, dependents } = this.info;
+
+    if (dependencies.length > 0)
+      rootElements.push({ title: "Partials", children: dependencies });
+
+    if (dependents.length > 0)
+      rootElements.push({ title: "References", children: dependents });
+
+    if (rootElements.length < 1) rootElements.push({ title: "None" });
+
+    return rootElements;
   }
 }
diff --git a/package.json b/package.json
index 620089e..1289cd6 100644
--- a/package.json
+++ b/package.json
@@ -7,7 +7,7 @@
   "author": "Ryan Paul",
   "publisher": "stripe",
   "license": "MIT",
-  "version": "0.0.11",
+  "version": "0.0.12",
   "description": "A Markdoc language server and Visual Studio Code extension",
   "repository": {
     "type": "git",
diff --git a/server/package.json b/server/package.json
index ecb427d..5a537e5 100644
--- a/server/package.json
+++ b/server/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@markdoc/language-server",
-  "version": "0.0.11",
+  "version": "0.0.12",
   "description": "A Markdoc language server",
   "main": "dist/index.js",
   "author": "Ryan Paul",