From 8e2991da1a813ad5181cc14eb93ac807d31a7ec8 Mon Sep 17 00:00:00 2001 From: Matt Low Date: Wed, 1 Jan 2025 05:06:37 +0000 Subject: [PATCH] Add pkg/util/dirtree Update dir_tree tool to use it Update filepicker bubble to use it --- pkg/agents/toolbox/dir_tree.go | 60 +---- pkg/tui/bubbles/filepicker/filepicker.go | 113 +++++++++ pkg/tui/bubbles/filepicker/model.go | 299 ----------------------- pkg/util/dirtree/render.go | 82 +++++++ pkg/util/dirtree/tree.go | 132 ++++++++++ 5 files changed, 340 insertions(+), 346 deletions(-) create mode 100644 pkg/tui/bubbles/filepicker/filepicker.go delete mode 100644 pkg/tui/bubbles/filepicker/model.go create mode 100644 pkg/util/dirtree/render.go create mode 100644 pkg/util/dirtree/tree.go diff --git a/pkg/agents/toolbox/dir_tree.go b/pkg/agents/toolbox/dir_tree.go index db8c926..2c953bd 100644 --- a/pkg/agents/toolbox/dir_tree.go +++ b/pkg/agents/toolbox/dir_tree.go @@ -3,12 +3,11 @@ package toolbox import ( "fmt" "os" - "path/filepath" "strconv" - "strings" toolutil "git.mlow.ca/mlow/lmcli/pkg/agents/toolbox/util" "git.mlow.ca/mlow/lmcli/pkg/api" + "git.mlow.ca/mlow/lmcli/pkg/util/dirtree" ) const TREE_DESCRIPTION = `Retrieve a tree-like view of a directory's contents. @@ -85,58 +84,25 @@ func tree(path string, depth int) api.CallResult { return api.CallResult{Message: reason} } - var treeOutput strings.Builder - treeOutput.WriteString(path + "\n") - err := buildTree(&treeOutput, path, "", depth) + tree := dirtree.NewTree(path) + err := tree.Root.LoadChildren(true, depth, nil) if err != nil { return api.CallResult{ Message: err.Error(), } } - return api.CallResult{Result: treeOutput.String()} -} - -func buildTree(output *strings.Builder, path string, prefix string, depth int) error { - files, err := os.ReadDir(path) - if err != nil { - return err - } - - for i, file := range files { - if strings.HasPrefix(file.Name(), ".") { - // Skip hidden files and directories - continue - } - - isLast := i == len(files)-1 - var branch string - if isLast { - branch = "└── " + nodes := dirtree.FlattenTree(tree.Root, true) + rendered := dirtree.RenderNodes(nodes, dirtree.Standard, -1, func(n *dirtree.Node) string { + if n.IsDir() { + return fmt.Sprintf("%s/", n.Name) } else { - branch = "├── " - } - - info, _ := file.Info() - size := info.Size() - sizeStr := fmt.Sprintf(" (%d bytes)", size) - - output.WriteString(prefix + branch + file.Name()) - if file.IsDir() { - output.WriteString("/\n") - if depth > 0 { - var nextPrefix string - if isLast { - nextPrefix = prefix + " " - } else { - nextPrefix = prefix + "│ " - } - buildTree(output, filepath.Join(path, file.Name()), nextPrefix, depth-1) + info, err := os.Stat(n.Path) + if err != nil { + return fmt.Sprintf("%s (ERR: %s bytes)", n.Name, err.Error()) } - } else { - output.WriteString(sizeStr + "\n") + return fmt.Sprintf("%s (%d bytes)", n.Name, info.Size()) } - } - - return nil + }) + return api.CallResult{Result: rendered} } diff --git a/pkg/tui/bubbles/filepicker/filepicker.go b/pkg/tui/bubbles/filepicker/filepicker.go new file mode 100644 index 0000000..45313c0 --- /dev/null +++ b/pkg/tui/bubbles/filepicker/filepicker.go @@ -0,0 +1,113 @@ +package filepicker + +import ( + "git.mlow.ca/mlow/lmcli/pkg/util/dirtree" + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" +) + +type Model struct { + tree *dirtree.Tree + cursor int + displayNodes []*dirtree.Node // Flattened view of visible nodes +} + +func New(path string) Model { + m := Model{ + tree: dirtree.NewTree(path), + } + + m.tree.Root.LoadChildren(false, 0, nil) + m.updateRenderTree() + return m +} + +func (m *Model) updateRenderTree() { + m.displayNodes = dirtree.FlattenTree(m.tree.Root, false) +} + +func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.KeyMsg: + switch msg.String() { + case "j": + if m.cursor < len(m.displayNodes)-1 { + m.cursor++ + } + case "k": + if m.cursor > 0 { + m.cursor-- + } + case "l", "right": + current := m.displayNodes[m.cursor] + if current.IsDir() { + current.Expanded = true + current.LoadChildren(false, 0, nil) + m.updateRenderTree() + } + case "h", "left": + current := m.displayNodes[m.cursor] + if current.IsDir() && current.Expanded { + current.Expanded = false + m.updateRenderTree() + } else { + // Find and collapse parent + if !current.IsRoot() { + current.Parent().Expanded = false + m.updateRenderTree() + // Move cursor to parent + for i, node := range m.displayNodes { + if node == current.Parent() { + m.cursor = i + break + } + } + } + } + case " ": + current := m.displayNodes[m.cursor] + if current.IsSelected() { + current.Deselect(true) + } else { + current.Select(true, 0) + } + m.updateRenderTree() + case "H": + m.tree.ShowHidden = !m.tree.ShowHidden + m.tree.Root.LoadChildren(false, 0, nil) + } + } + return m, nil +} + +func (m Model) Init() tea.Cmd { + return nil +} + +func (m Model) View() string { + // Define styles + dirStyle := lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("12")) // Blue for directories + fileStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("7")) // White for files + selectedStyle := lipgloss.NewStyle().Bold(true) + + renderNode := func(node *dirtree.Node) string { + var nodeName string + if node.IsDir() { + caret := "▶ " + if node.Expanded { + caret = "▼ " + } + nodeName = caret + node.Name + "/" + nodeName = dirStyle.Render(nodeName) + } else { + nodeName = fileStyle.Render(node.Name) + } + + if node.IsSelected() { + nodeName = selectedStyle.Render(nodeName) + } + return nodeName + } + + return dirtree.RenderNodes(m.displayNodes, dirtree.Compact, m.cursor, renderNode) +} diff --git a/pkg/tui/bubbles/filepicker/model.go b/pkg/tui/bubbles/filepicker/model.go deleted file mode 100644 index d0835f3..0000000 --- a/pkg/tui/bubbles/filepicker/model.go +++ /dev/null @@ -1,299 +0,0 @@ -package filepicker - -import ( - "fmt" - "os" - "path/filepath" - "strings" - - tea "github.com/charmbracelet/bubbletea" - "github.com/charmbracelet/lipgloss" -) - -type Node struct { - path string - name string - isDir bool - expanded bool - children []*Node -} - -type Model struct { - root *Node - cursor int - showHidden bool - displayLines []*Node // Flattened view of visible nodes - rootPath string - - SelectedPaths []string - expandedPaths map[string]bool // Track expanded paths -} - -func New(path string) Model { - root := &Node{ - path: path, - name: filepath.Base(path), - isDir: true, - expanded: true, - } - - m := Model{ - root: root, - rootPath: path, - SelectedPaths: make([]string, 0), - showHidden: false, - expandedPaths: make(map[string]bool), - } - - m.loadDirectory(root) - m.updateDisplayLines() - return m -} - -func (m *Model) loadDirectory(node *Node) error { - // Save current expanded state of children - for _, child := range node.children { - if child.expanded { - m.expandedPaths[child.path] = true - } - } - - entries, err := os.ReadDir(node.path) - if err != nil { - return err - } - - node.children = make([]*Node, 0) - for _, entry := range entries { - if !m.showHidden && strings.HasPrefix(entry.Name(), ".") { - continue - } - - childPath := filepath.Join(node.path, entry.Name()) - child := &Node{ - path: childPath, - name: entry.Name(), - isDir: entry.IsDir(), - expanded: m.expandedPaths[childPath], // Restore expanded state - } - node.children = append(node.children, child) - } - return nil -} - -func (m *Model) updateDisplayLines() { - m.displayLines = make([]*Node, 0) - m.flattenTree(m.root, 0) -} - -func (m *Model) flattenTree(node *Node, depth int) { - m.displayLines = append(m.displayLines, node) - if node.isDir && node.expanded { - for _, child := range node.children { - m.flattenTree(child, depth+1) - } - } -} - -func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) { - switch msg := msg.(type) { - case tea.KeyMsg: - switch msg.String() { - case "j": - if m.cursor < len(m.displayLines)-1 { - m.cursor++ - } - case "k": - if m.cursor > 0 { - m.cursor-- - } - case "l": - current := m.displayLines[m.cursor] - if current.isDir { - current.expanded = true - m.expandedPaths[current.path] = true - m.loadDirectory(current) - m.updateDisplayLines() - } - case "h": - current := m.displayLines[m.cursor] - if current.isDir && current.expanded { - current.expanded = false - delete(m.expandedPaths, current.path) - m.updateDisplayLines() - } else { - // Find and collapse parent - parent := m.findParent(current) - if parent != nil && parent.path != m.rootPath { - parent.expanded = false - delete(m.expandedPaths, parent.path) - m.updateDisplayLines() - // Move cursor to parent - for i, node := range m.displayLines { - if node == parent { - m.cursor = i - break - } - } - } - } - case " ": - current := m.displayLines[m.cursor] - if m.isSelected(current.path) { - m.SelectedPaths = removeString(m.SelectedPaths, current.path) - } else { - m.SelectedPaths = append(m.SelectedPaths, current.path) - } - case "H": - m.showHidden = !m.showHidden - m.reloadTree() - } - } - return m, nil -} - -func (m Model) Init() tea.Cmd { - return nil -} - -func (m Model) View() string { - var sb strings.Builder - - // Define styles - dirStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("12")) // Blue for directories - fileStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("7")) // White for files - cursorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("yellow")) - selectedStyle := lipgloss.NewStyle().Bold(true) - - // Track the parent-child relationships for proper indentation - levelIsLast := make(map[int]bool) - - for i, node := range m.displayLines { - // Calculate depth and prefix - depth := 0 - if node != m.root { - depth = strings.Count(node.path[len(m.rootPath):], string(os.PathSeparator)) - } - - // Build the prefix (tree structure) - var prefix strings.Builder - - if node != m.root { - // For each level up to current depth, add the appropriate prefix - for d := 0; d < depth; d++ { - if levelIsLast[d] { - prefix.WriteString(" ") - } else { - prefix.WriteString("│ ") - } - } - - // Find if this node is the last among its siblings - parent := m.findParent(node) - if parent != nil { - isLast := parent.children[len(parent.children)-1] == node - levelIsLast[depth] = isLast - - if isLast { - prefix.WriteString("└── ") - } else { - prefix.WriteString("├── ") - } - } - } - - // Cursor indicator - cursor := " " - if i == m.cursor { - cursor = cursorStyle.Render(">") - } - - // Node name with optional caret for directories - var nodeName string - if node.isDir { - caret := "▶ " - if node.expanded { - caret = "▼ " - } - nodeName = caret + node.name - if node != m.root { - nodeName += "/" - } - nodeName = dirStyle.Render(nodeName) - } else { - nodeName = fileStyle.Render(node.name) - } - - // Apply bold style if selected - if m.isSelected(node.path) { - nodeName = selectedStyle.Render(nodeName) - } - - // Build the line - line := fmt.Sprintf("%s %s%s\n", - cursor, - prefix.String(), - nodeName) - - sb.WriteString(line) - } - - return sb.String() -} - -func (m *Model) reloadTree() { - // Save expanded state before reloading - m.saveExpandedState(m.root) - - m.root.children = nil - m.loadDirectory(m.root) - m.updateDisplayLines() -} - -func (m *Model) saveExpandedState(node *Node) { - if node.isDir && node.expanded { - m.expandedPaths[node.path] = true - for _, child := range node.children { - m.saveExpandedState(child) - } - } -} - -func (m *Model) findParent(node *Node) *Node { - var findParentRec func(*Node) *Node - - findParentRec = func(current *Node) *Node { - if current == nil { - return nil - } - for _, child := range current.children { - if child == node { - return current - } - if result := findParentRec(child); result != nil { - return result - } - } - return nil - } - - return findParentRec(m.root) -} - -func (m Model) isSelected(path string) bool { - for _, selected := range m.SelectedPaths { - if selected == path { - return true - } - } - return false -} - -func removeString(slice []string, s string) []string { - for i, v := range slice { - if v == s { - return append(slice[:i], slice[i+1:]...) - } - } - return slice -} diff --git a/pkg/util/dirtree/render.go b/pkg/util/dirtree/render.go new file mode 100644 index 0000000..1104c30 --- /dev/null +++ b/pkg/util/dirtree/render.go @@ -0,0 +1,82 @@ +package dirtree + +import "strings" + +// FlattenTree returns []*Node consistent with each node's `Expanded` state +func FlattenTree(node *Node, expandAll bool) []*Node { + lines := make([]*Node, 0) + var flatten func(n *Node) + flatten = func(node *Node) { + lines = append(lines, node) + if node.isDir && (node.Expanded || expandAll) { + for _, child := range node.children { + flatten(child) + } + } + } + flatten(node) + return lines +} + +type TreeFormat struct { + Space string + Span string + Node string + NodeLast string +} + +var ( + Standard = TreeFormat{ + Space: " ", + Span: "│ ", + Node: "├── ", + NodeLast: "└── ", + } + Compact = TreeFormat{ + Space: " ", + Span: "│ ", + Node: "├ ", + NodeLast: "└ ", + } +) + +func RenderNodes(displayNodes []*Node, format TreeFormat, cursor int, renderNode func(*Node) string) string { + var sb strings.Builder + levelIsLast := make(map[int]bool) + for i, node := range displayNodes { + // Cursor indicator + if cursor >= 0 { + if i == cursor { + sb.WriteString("> ") + } else { + sb.WriteString(" ") + } + } + + // Tree structure + if node.depth > 0 { + for d := 1; d < node.depth; d++ { + if levelIsLast[d] { + sb.WriteString(format.Space) + } else { + sb.WriteString(format.Span) + } + } + + // Find if this node is the last among its siblings + isLast := node.parent.children[len(node.parent.children)-1] == node + levelIsLast[node.depth] = isLast + + if isLast { + sb.WriteString(format.NodeLast) + } else { + sb.WriteString(format.Node) + } + } + + // Render the node + sb.WriteString(renderNode(node)) + sb.WriteString("\n") + } + return sb.String() +} diff --git a/pkg/util/dirtree/tree.go b/pkg/util/dirtree/tree.go new file mode 100644 index 0000000..cbf4c26 --- /dev/null +++ b/pkg/util/dirtree/tree.go @@ -0,0 +1,132 @@ +package dirtree + +import ( + "os" + "path/filepath" + "strings" +) + +type Tree struct { + Root *Node + ShowHidden bool +} + +type Node struct { + tree *Tree + parent *Node + children []*Node + depth int + Path string + Name string + Expanded bool + isDir bool + selected bool +} + +func NewTree(path string) *Tree { + tree := &Tree{ + ShowHidden: false, + } + tree.Root = &Node{ + tree: tree, + parent: nil, + Path: path, + Name: filepath.Base(path), + isDir: true, + Expanded: true, + } + return tree +} + +func (node *Node) IsRoot() bool { + return node.tree.Root == node +} + +func (node *Node) Parent() *Node { + return node.parent +} + +func (node *Node) IsSelected() bool { + return node.selected +} + +func (node *Node) IsDir() bool { + return node.isDir +} + +func (node *Node) Select(recurse bool, maxDepth int) { + node.selected = true + if node.isDir && recurse { + node.LoadChildren(recurse, maxDepth, func (n *Node) { + n.selected = true + }) + } +} + +func (node *Node) Deselect(recurse bool) { + node.selected = false + if node.isDir && recurse { + for _, child := range node.children { + child.Deselect(recurse) + } + } +} + +func (node *Node) LoadChildren(recurse bool, maxDepth int, visitFunc func (node *Node)) error { + startDepth := node.depth + var loadChildren func(node *Node) error + loadChildren = func(node *Node) error { + // Load new entries + entries, err := os.ReadDir(node.Path) + if err != nil { + return err + } + + // Save current children before loading + oldChildren := make(map[string]*Node) + for _, child := range node.children { + oldChildren[child.Path] = child + } + + node.children = make([]*Node, 0) + for _, entry := range entries { + if !node.tree.ShowHidden && strings.HasPrefix(entry.Name(), ".") { + continue + } + + child := &Node{ + tree: node.tree, + parent: node, + depth: node.depth + 1, + Name: entry.Name(), + isDir: entry.IsDir(), + Path: filepath.Join(node.Path, entry.Name()), + } + + oldChild, exists := oldChildren[child.Path] + if child.isDir { + // Preserve child's grandchildren and expanded/selected state + if exists && (oldChild.Expanded || len(oldChild.children) > 0) { + child.Expanded = oldChild.Expanded + child.selected = oldChild.selected + child.children = oldChild.children + for i := range child.children { + child.children[i].parent = child + } + loadChildren(child) + } else if recurse && (maxDepth <= 0 || node.depth - startDepth < maxDepth) { + loadChildren(child) + } + } else if exists { + child.selected = oldChild.selected + } + + node.children = append(node.children, child) + if visitFunc != nil { + visitFunc(child) + } + } + return nil + } + return loadChildren(node) +}