From c4d9ee4bc0465ceaaef6dd170d9eeb0ad3ab75e0 Mon Sep 17 00:00:00 2001
From: Tony Crisci <tony@dubstepdish.com>
Date: Thu, 14 Sep 2017 18:05:09 -0400
Subject: [PATCH] Add descendents iterator to Con

You can now iterate through the Con to get a generator of all the
descendents.

fixes #66
---
 README.rst     |  6 ++++++
 i3ipc/i3ipc.py | 48 ++++++++++++++++++++++++++----------------------
 2 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/README.rst b/README.rst
index 971a69b..ace22af 100644
--- a/README.rst
+++ b/README.rst
@@ -65,6 +65,12 @@ Example
     for container in i3.get_tree().find_fullscreen():
         container.command('fullscreen')
 
+    # Print the names of all the containers in the tree
+    root = i3.get_tree()
+    print(root.name)
+    for con in root:
+        print(con.name)
+
     # Define a callback to be called when you switch workspaces.
     def on_workspace_focus(self, e):
         # The first parameter is the connection to the ipc and the second is an object
diff --git a/i3ipc/i3ipc.py b/i3ipc/i3ipc.py
index 27d7628..8f7cc62 100644
--- a/i3ipc/i3ipc.py
+++ b/i3ipc/i3ipc.py
@@ -8,6 +8,7 @@
 import re
 import subprocess
 from enum import Enum
+from collections import deque
 
 
 class MessageType(Enum):
@@ -875,6 +876,20 @@ def __init__(self, data, parent, conn):
         if 'gaps' in data:
             self.gaps = Gaps(data['gaps'])
 
+    def __iter__(self):
+        """
+        Iterate through the descendents of this node (breadth-first tree traversal)
+        """
+        queue = deque()
+        queue.append(self)
+
+        while queue:
+            con = queue.popleft()
+            if not con is self:
+                yield con
+
+            queue.extend(con.nodes)
+
     def root(self):
         """
         Retrieves the root container.
@@ -899,18 +914,7 @@ def descendents(self):
 
         :rtype: List of :class:`Con`.
         """
-        descendents = []
-
-        def collect_descendents(con):
-            for c in con.nodes:
-                descendents.append(c)
-                collect_descendents(c)
-            for c in con.floating_nodes:
-                descendents.append(c)
-                collect_descendents(c)
-
-        collect_descendents(self)
-        return descendents
+        return [c for c in self]
 
     def leaves(self):
         """
@@ -922,7 +926,7 @@ def leaves(self):
         """
         leaves = []
 
-        for c in self.descendents():
+        for c in self:
             if not c.nodes and c.type == "con" and c.parent.type != "dockarea":
                 leaves.append(c)
 
@@ -978,45 +982,45 @@ def find_focused(self):
         :rtype class Con:
         """
         try:
-            return next(c for c in self.descendents() if c.focused)
+            return next(c for c in self if c.focused)
         except StopIteration:
             return None
 
     def find_by_id(self, id):
         try:
-            return next(c for c in self.descendents() if c.id == id)
+            return next(c for c in self if c.id == id)
         except StopIteration:
             return None
 
     def find_by_window(self, window):
         try:
-            return next(c for c in self.descendents() if c.window == window)
+            return next(c for c in self if c.window == window)
         except StopIteration:
             return None
 
     def find_by_role(self, pattern):
-        return [c for c in self.descendents()
+        return [c for c in self
                 if c.window_role and re.search(pattern, c.window_role)]
 
     def find_named(self, pattern):
-        return [c for c in self.descendents()
+        return [c for c in self
                 if c.name and re.search(pattern, c.name)]
 
     def find_classed(self, pattern):
-        return [c for c in self.descendents()
+        return [c for c in self
                 if c.window_class and re.search(pattern, c.window_class)]
 
     def find_instanced(self, pattern):
-        return [c for c in self.descendents()
+        return [c for c in self
                 if c.window_instance and re.search(pattern, c.window_instance)]
 
     def find_marked(self, pattern=".*"):
         pattern = re.compile(pattern)
-        return [c for c in self.descendents()
+        return [c for c in self
                 if any(pattern.search(mark) for mark in c.marks)]
 
     def find_fullscreen(self):
-        return [c for c in self.descendents()
+        return [c for c in self
                 if c.type == 'con' and c.fullscreen_mode]
 
     def workspace(self):