feat: add charset option

treelike.go

@@ -54,6 +54,10 @@ func getOpts() Options {
 		case "-c", "--charset":
 			{
 				opts.charset = args[1]
+				if opts.charset != "utf-8" && opts.charset != "ascii" {
+					fmt.Fprintf(os.Stderr, "Invalid charset: %s\n", opts.charset)
+					os.Exit(1)
+				}
 				args = args[2:]
 			}
 		case "-s", "--trailing-slash":
@@ -133,7 +137,7 @@ func parseInput(input string) *Node {
 }
 
 func describeTree(node *Node, opts *Options) string {
-	lines := []string{getAsciiLine(node, opts)}
+	lines := []string{getTreeLine(node, opts)}
 
 	for _, child := range node.children {
 		next := describeTree(child, opts)
@@ -155,13 +159,27 @@ func describeTree(node *Node, opts *Options) string {
 }
 
 const (
-	CHILD      string = "├── "
-	LAST_CHILD string = "└── "
-	DIRECTORY  string = "│   "
-	EMPTY      string = "    "
+	UTF8_CHILD      string = "├── "
+	UTF8_LAST_CHILD string = "└── "
+	UTF8_DIRECTORY  string = "│   "
+	UTF8_EMPTY      string = "    "
 )
 
-func getAsciiLine(node *Node, opts *Options) string {
+const (
+	ASCII_CHILD      string = "|-- "
+	ASCII_LAST_CHILD string = "`-- "
+	ASCII_DIRECTORY  string = "|   "
+	ASCII_EMPTY      string = "    "
+)
+
+func getPrefixes(opts *Options) (string, string, string, string) {
+	if opts.charset == "ascii" {
+		return ASCII_CHILD, ASCII_LAST_CHILD, ASCII_DIRECTORY, ASCII_EMPTY
+	}
+	return UTF8_CHILD, UTF8_LAST_CHILD, UTF8_DIRECTORY, UTF8_EMPTY
+}
+
+func getTreeLine(node *Node, opts *Options) string {
 	if node.parent == nil {
 		if opts.rootDot {
 			return node.name
@@ -170,6 +188,8 @@ func getAsciiLine(node *Node, opts *Options) string {
 		}
 	}
 
+	CHILD, LAST_CHILD, DIRECTORY, EMPTY := getPrefixes(opts)
+
 	var chunks strings.Builder
 
 	if isLastChild(node) {
@@ -194,7 +214,7 @@ func getAsciiLine(node *Node, opts *Options) string {
 	if opts.rootDot {
 		return str
 	}
-	return removePrefix(str)
+	return removePrefix(str, opts)
 }
 
 func getName(node *Node, opts *Options) string {
@@ -220,7 +240,9 @@ func isLastChild(node *Node) bool {
 	return node.parent != nil && node.parent.children[len(node.parent.children)-1] == node
 }
 
-func removePrefix(str string) string {
+func removePrefix(str string, opts *Options) string {
+	CHILD, LAST_CHILD, DIRECTORY, EMPTY := getPrefixes(opts)
+
 	prefixes := []string{CHILD, LAST_CHILD, DIRECTORY, EMPTY}
 
 	for _, prefix := range prefixes {

treelike_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"strings"
 	"testing"
 )
 
@@ -69,10 +70,10 @@ func TestGetAsciiLine(t *testing.T) {
 	root.children = append(root.children, child)
 	opts := Options{rootDot: true}
 
-	result := getAsciiLine(child, &opts)
+	result := getTreeLine(child, &opts)
 	expected := "└── child"
 	if result != expected {
-		t.Errorf("getAsciiLine() = %q; want %q", result, expected)
+		t.Errorf("getTreeLine() = %q; want %q", result, expected)
 	}
 }
 
@@ -110,6 +111,9 @@ func TestIsLastChild(t *testing.T) {
 }
 
 func TestRemovePrefix(t *testing.T) {
+	opts := Options{false, "", strings.Builder{}, "utf-8", false, false, true}
+	CHILD, LAST_CHILD, DIRECTORY, EMPTY := getPrefixes(&opts)
+
 	tests := []struct {
 		str      string
 		expected string
@@ -121,9 +125,38 @@ func TestRemovePrefix(t *testing.T) {
 	}
 
 	for _, test := range tests {
-		result := removePrefix(test.str)
+		result := removePrefix(test.str, &opts)
 		if result != test.expected {
 			t.Errorf("removePrefix(%q) = %q; want %q", test.str, result, test.expected)
 		}
 	}
 }
+func TestGetPrefixes(t *testing.T) {
+	tests := []struct {
+		charset           string
+		expectedChild     string
+		expectedLastChild string
+		expectedDirectory string
+		expectedEmpty     string
+	}{
+		{"utf-8", UTF8_CHILD, UTF8_LAST_CHILD, UTF8_DIRECTORY, UTF8_EMPTY},
+		{"ascii", ASCII_CHILD, ASCII_LAST_CHILD, ASCII_DIRECTORY, ASCII_EMPTY},
+	}
+
+	for _, test := range tests {
+		opts := &Options{charset: test.charset}
+		child, lastChild, directory, empty := getPrefixes(opts)
+		if child != test.expectedChild {
+			t.Errorf("For charset %s, expected child prefix %s, but got %s", test.charset, test.expectedChild, child)
+		}
+		if lastChild != test.expectedLastChild {
+			t.Errorf("For charset %s, expected last child prefix %s, but got %s", test.charset, test.expectedLastChild, lastChild)
+		}
+		if directory != test.expectedDirectory {
+			t.Errorf("For charset %s, expected directory prefix %s, but got %s", test.charset, test.expectedDirectory, directory)
+		}
+		if empty != test.expectedEmpty {
+			t.Errorf("For charset %s, expected empty prefix %s, but got %s", test.charset, test.expectedEmpty, empty)
+		}
+	}
+}