From d4e18dd66f48a258334ad82b50dd5610cf2d8896 Mon Sep 17 00:00:00 2001
From: Anton Medvedev <anton@medv.io>
Date: Sun, 10 Mar 2024 20:48:46 +0100
Subject: [PATCH] Add 1-9 keys to expand ap to nth level

---
 help.go      |  6 +++++-
 keymap.go    |  5 +++++
 main.go      | 14 ++++++++++++--
 node.go      | 20 +++++++++++++++-----
 node_test.go |  9 +++++++++
 5 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/help.go b/help.go
index 6ed8753..09de177 100644
--- a/help.go
+++ b/help.go
@@ -53,7 +53,11 @@ func keyMapInfo(keyMap KeyMap, style lipgloss.Style) []string {
 		k := v.Field(i).Interface().(key.Binding)
 		str := k.Help().Key
 		if len(str) == 0 {
-			str = strings.Join(k.Keys(), ", ")
+			if len(k.Keys()) > 5 {
+				str = fmt.Sprintf("%v-%v", k.Keys()[0], k.Keys()[len(k.Keys())-1])
+			} else {
+				str = strings.Join(k.Keys(), ", ")
+			}
 		}
 		keys = append(keys, fmt.Sprintf("%v    ", str))
 	}
diff --git a/keymap.go b/keymap.go
index 4342c7a..2ca7311 100644
--- a/keymap.go
+++ b/keymap.go
@@ -19,6 +19,7 @@ type KeyMap struct {
 	CollapseRecursively key.Binding
 	ExpandAll           key.Binding
 	CollapseAll         key.Binding
+	CollapseLevel       key.Binding
 	NextSibling         key.Binding
 	PrevSibling         key.Binding
 	ToggleWrap          key.Binding
@@ -97,6 +98,10 @@ func init() {
 			key.WithKeys("E"),
 			key.WithHelp("", "collapse all"),
 		),
+		CollapseLevel: key.NewBinding(
+			key.WithKeys("1", "2", "3", "4", "5", "6", "7", "8", "9"),
+			key.WithHelp("", "collapse to nth level"),
+		),
 		NextSibling: key.NewBinding(
 			key.WithKeys("J", "shift+down"),
 			key.WithHelp("", "next sibling"),
diff --git a/main.go b/main.go
index b29bcf0..026a127 100644
--- a/main.go
+++ b/main.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
+	"math"
 	"os"
 	"path"
 	"regexp"
@@ -485,7 +486,7 @@ func (m *model) handleKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
 	case key.Matches(msg, keyMap.ExpandRecursively):
 		n := m.cursorPointsTo()
 		if n.hasChildren() {
-			n.expandRecursively()
+			n.expandRecursively(0, math.MaxInt)
 		}
 		m.showCursor = true
 
@@ -507,7 +508,7 @@ func (m *model) handleKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
 		at := m.cursorPointsTo()
 		n := m.top
 		for n != nil {
-			n.expandRecursively()
+			n.expandRecursively(0, math.MaxInt)
 			if n.end == nil {
 				n = nil
 			} else {
@@ -516,6 +517,15 @@ func (m *model) handleKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
 		}
 		m.selectNode(at)
 
+	case key.Matches(msg, keyMap.CollapseLevel):
+		at := m.cursorPointsTo()
+		if at != nil && at.hasChildren() {
+			toLevel, _ := strconv.Atoi(msg.String())
+			at.collapseRecursively()
+			at.expandRecursively(0, toLevel)
+			m.showCursor = true
+		}
+
 	case key.Matches(msg, keyMap.ToggleWrap):
 		at := m.cursorPointsTo()
 		m.wrap = !m.wrap
diff --git a/node.go b/node.go
index ac5eca4..efbf58d 100644
--- a/node.go
+++ b/node.go
@@ -137,11 +137,21 @@ func (n *node) expand() {
 	}
 }
 
-func (n *node) expandRecursively() {
-	at := n
-	for at != nil && at != n.end {
-		at.expand()
-		at = at.next
+func (n *node) expandRecursively(level, maxLevel int) {
+	if level >= maxLevel {
+		return
+	}
+	if n.isCollapsed() {
+		n.expand()
+	}
+	it := n.next
+	for it != nil && it != n.end {
+		if it.hasChildren() {
+			it.expandRecursively(level+1, maxLevel)
+			it = it.end.next
+		} else {
+			it = it.next
+		}
 	}
 }
 
diff --git a/node_test.go b/node_test.go
index 6bfb30f..9f37edd 100644
--- a/node_test.go
+++ b/node_test.go
@@ -24,3 +24,12 @@ func TestNode_children(t *testing.T) {
 	paths, _ := n.children()
 	assert.Equal(t, []string{"a", "b", "c"}, paths)
 }
+
+func TestNode_expandRecursively(t *testing.T) {
+	n, err := parse([]byte(`{"a": {"b": {"c": 1}}}`))
+	require.NoError(t, err)
+
+	n.collapseRecursively()
+	n.expandRecursively(0, 3)
+	assert.Equal(t, `"c"`, string(n.next.next.next.key))
+}