Refactor the output.go code to use official Tree Sitter go bindings

spachava753 · Oct 9, 2024 · a21d8f1 · a21d8f1
1 parent bb08069
commit a21d8f1
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 20 deletions.
diff --git a/codemap/output.go b/codemap/output.go
@@ -67,7 +67,7 @@ func generateFileOutput(fsys fs.FS, path string, maxLiteralLen int) (string, err
 	// Traverse the tree and extract relevant information
 	var traverse func(node *sitter.Node)
 	traverse = func(node *sitter.Node) {
-		switch node.GrammarName() {
+		switch node.Kind() {
 		case "source_file":
 			for i := 0; i < int(node.ChildCount()); i++ {
 				traverse(node.Child(uint(i)))
@@ -78,7 +78,7 @@ func generateFileOutput(fsys fs.FS, path string, maxLiteralLen int) (string, err
 			for i := 0; i < int(node.NamedChildCount()); i++ {
 				child := node.NamedChild(uint(i))
 				cName := node.FieldNameForChild(uint32(i + 1))
-				if child.GrammarName() != "block" {
+				if child.Kind() != "block" {
 					output.WriteString(child.Utf8Text(src))
 				}
 				switch cName {
@@ -92,8 +92,12 @@ func generateFileOutput(fsys fs.FS, path string, maxLiteralLen int) (string, err
 		case "comment":
 			output.WriteString(node.Utf8Text(src))
 			output.WriteString("\n")
+		case "type_declaration":
+			output.WriteString(node.Utf8Text(src))
+			output.WriteString("\n\n")
 		default:
 			output.WriteString(node.Utf8Text(src))
+			output.WriteString("\n")
 		}
 	}
 
@@ -113,9 +117,9 @@ func traverseAndTruncate(node *sitter.Node, src []byte, output *strings.Builder,
 		child := node.Child(uint(i))
 		cchild := child.Utf8Text(src)
 		_ = cchild
-		tchild := child.GrammarName()
+		tchild := child.Kind()
 		_ = tchild
-		switch child.GrammarName() {
+		switch child.Kind() {
 		case "interpreted_string_literal", "raw_string_literal":
 			content := child.Utf8Text(src)
 			if len(content)-2 > maxLiteralLen {
@@ -126,10 +130,17 @@ func traverseAndTruncate(node *sitter.Node, src []byte, output *strings.Builder,
 			}
 		case "(":
 			output.WriteString(child.Utf8Text(src))
-			if node.GrammarName() == "var_spec_list" ||
-				(node.GrammarName() == "const_declaration" && node.NamedChildCount() > 0) {
+			if node.Kind() == "var_spec_list" ||
+				(node.Kind() == "const_declaration" && node.NamedChildCount() > 0) {
 				output.WriteString("\n")
 			}
+		case ")":
+			if node.Kind() == "var_spec_list" ||
+				(node.Kind() == "const_declaration" && node.NamedChildCount() > 0) {
+				output.WriteString("\n)\n")
+			} else {
+				output.WriteString(")")
+			}
 		case "=":
 			output.WriteString(child.Utf8Text(src))
 			output.WriteString(" ")
@@ -138,8 +149,8 @@ func traverseAndTruncate(node *sitter.Node, src []byte, output *strings.Builder,
 			output.WriteString(" ")
 		case "identifier":
 			for p := node; p != nil; p = p.Parent() {
-				if p.GrammarName() == "var_spec_list" ||
-					(p.GrammarName() == "const_declaration" && p.NamedChildCount() > 1) {
+				if p.Kind() == "var_spec_list" ||
+					(p.Kind() == "const_declaration" && p.NamedChildCount() > 1) {
 					output.WriteString("\t")
 					break
 				}
@@ -148,8 +159,8 @@ func traverseAndTruncate(node *sitter.Node, src []byte, output *strings.Builder,
 			output.WriteString(" ")
 		case "comment":
 			for p := node; p != nil; p = p.Parent() {
-				if p.GrammarName() == "var_spec_list" ||
-					(p.GrammarName() == "const_declaration" && p.NamedChildCount() > 1) {
+				if p.Kind() == "var_spec_list" ||
+					(p.Kind() == "const_declaration" && p.NamedChildCount() > 1) {
 					output.WriteString("\t")
 					break
 				}
@@ -158,6 +169,11 @@ func traverseAndTruncate(node *sitter.Node, src []byte, output *strings.Builder,
 			if child.PrevNamedSibling() == nil || child.PrevNamedSibling().EndPosition().Row != child.EndPosition().Row {
 				output.WriteString("\n")
 			}
+		case "var_spec", "const_spec":
+			traverseAndTruncate(child, src, output, maxLiteralLen)
+			if child.NextNamedSibling() == nil || child.NextNamedSibling().StartPosition().Row > child.StartPosition().Row {
+				output.WriteString("\n")
+			}
 		default:
 			if child.ChildCount() > 0 {
 				traverseAndTruncate(child, src, output, maxLiteralLen)

diff --git a/codemap/output_test.go b/codemap/output_test.go
@@ -237,17 +237,12 @@ var (
 package globals
 
 var (
-	LongString = "T..."
-
-	LongStringBackTick = ` + "`" + `T...` + "`" + `
-
-	LongByteSlice = []byte("T...")
-
+	LongString            = "T..."
+	LongStringBackTick    = ` + "`" + `T...` + "`" + `
+	LongByteSlice         = []byte("T...")
 	LongByteSliceBackTick = []byte(` + "`" + `T...` + "`" + `)
-
-	ShortString = "T"
-
-	RegularVar = 42
+	ShortString           = "T"
+	RegularVar            = 42
 )
 </file_map>
 </file>