From 470ebafea46471a5faf7e862f9587e32d782e19e Mon Sep 17 00:00:00 2001
From: smoka7 <36933074+smoka7@users.noreply.github.com>
Date: Fri, 11 Feb 2022 22:41:31 +0330
Subject: [PATCH] add a function to traverse in nodes (#3)

---
 sway_test.go | 14 ++++++++++++++
 types.go     | 10 +++++++++-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/sway_test.go b/sway_test.go
index 7d3e39e..272fa48 100644
--- a/sway_test.go
+++ b/sway_test.go
@@ -148,3 +148,17 @@ func processFocus(ctx context.Context, client sway.Client, node *sway.Node) {
 		log.Println(err)
 	}
 }
+
+func TestFocused(t *testing.T) {
+	ctx := context.Background()
+	client, err := sway.New(ctx)
+	if err != nil {
+		t.Error(err)
+	}
+	tree, err := client.GetTree(ctx)
+	if err != nil {
+		t.Error(err)
+	}
+	f := tree.FocusedNode()
+	printJSON(f)
+}
diff --git a/types.go b/types.go
index ee8cd3e..78d691c 100644
--- a/types.go
+++ b/types.go
@@ -124,6 +124,14 @@ type Node struct {
 
 // FocusedNode traverses the node tree and returns the focused node
 func (n *Node) FocusedNode() *Node {
+	focusedNode := n.TraverseNodes(func(n *Node) bool {
+		return n.Focused
+	})
+	return focusedNode
+}
+
+// TraverseNodes returns the first Node matching the predicate
+func (n *Node) TraverseNodes(predicate func(*Node) bool) *Node {
 	queue := []*Node{n}
 	for len(queue) > 0 {
 		n = queue[0]
@@ -133,7 +141,7 @@ func (n *Node) FocusedNode() *Node {
 			continue
 		}
 
-		if n.Focused {
+		if predicate(n) {
 			return n
 		}