From 32e7f4ab199809988d577328d0a3842dc850a299 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Wed, 1 Jan 2025 21:00:45 -0700
Subject: [PATCH] fix classes keyword

---
 d2lsp/completion.go      |  3 +++
 d2lsp/completion_test.go | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/d2lsp/completion.go b/d2lsp/completion.go
index d86c98570..86e683ce7 100644
--- a/d2lsp/completion.go
+++ b/d2lsp/completion.go
@@ -145,6 +145,9 @@ func getKeywordContext(text string, m *d2ast.Map, line, column int) string {
 		if _, isBoard := d2ast.BoardKeywords[firstPart]; isBoard {
 			firstPart = ""
 		}
+		if firstPart == "classes" {
+			firstPart = ""
+		}
 
 		_, isHolder := d2ast.ReservedKeywordHolders[firstPart]
 		if !isHolder {
diff --git a/d2lsp/completion_test.go b/d2lsp/completion_test.go
index b8aa60b60..8a8ee901e 100644
--- a/d2lsp/completion_test.go
+++ b/d2lsp/completion_test.go
@@ -30,6 +30,18 @@ func TestGetCompletionItems(t *testing.T) {
 			column: 8,
 			want:   getStyleCompletions(),
 		},
+		{
+			name: "classes shapes",
+			text: `classes: {
+  goal: {
+    shape:
+  }
+}
+`,
+			line:   2,
+			column: 10,
+			want:   getShapeCompletions(),
+		},
 		{
 			name: "nested style map suggestions",
 			text: `a: {
@@ -291,6 +303,15 @@ layers: {
 			column: 16,
 			want:   getShapeCompletions(),
 		},
+		{
+			name: "shape 2 suggestions",
+			text: `a: {
+  shape:
+}`,
+			line:   1,
+			column: 8,
+			want:   getShapeCompletions(),
+		},
 	}
 
 	for _, tt := range tests {