From f604b73a3c9f4afea3638951bb7677d459648c1d Mon Sep 17 00:00:00 2001
From: Alexey Shrub <ashrub@yandex.ru>
Date: Thu, 25 Oct 2018 11:08:03 +0300
Subject: [PATCH] Tree walk tests

---
 lark/tree.py        |  4 ----
 tests/test_trees.py | 11 +++++++++++
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/lark/tree.py b/lark/tree.py
index b2c7a60..e892a44 100644
--- a/lark/tree.py
+++ b/lark/tree.py
@@ -102,10 +102,6 @@ class Tree(object):
                 seen.add(id(x))
 
     def iter_subtrees_topdown(self):
-        if not self.children:
-            yield self
-            return
-
         stack = [self]
         while stack:
             node = stack.pop()
diff --git a/tests/test_trees.py b/tests/test_trees.py
index 564d02b..7e1d841 100644
--- a/tests/test_trees.py
+++ b/tests/test_trees.py
@@ -21,6 +21,17 @@ class TestTrees(TestCase):
         data = pickle.dumps(s)
         assert pickle.loads(data) == s
 
+    def test_iter_subtrees(self):
+        expected = [Tree('b', 'x'), Tree('c', 'y'), Tree('d', 'z'),
+                    Tree('a', [Tree('b', 'x'), Tree('c', 'y'), Tree('d', 'z')])]
+        nodes = list(self.tree1.iter_subtrees())
+        self.assertEqual(nodes, expected)
+
+    def test_iter_subtrees_topdown(self):
+        expected = [Tree('a', [Tree('b', 'x'), Tree('c', 'y'), Tree('d', 'z')]),
+                    Tree('b', 'x'), Tree('c', 'y'), Tree('d', 'z')]
+        nodes = list(self.tree1.iter_subtrees_topdown())
+        self.assertEqual(nodes, expected)
 
     def test_interp(self):
         t = Tree('a', [Tree('b', []), Tree('c', []), 'd'])