Add a generic node iterator

Suggested by Mitja Horvat (pinkfluid) in
https://github.com/ulfalizer/Kconfiglib/pull/50.

Kconfig.node_iter() iterates through all menu nodes in the menu tree in
Kconfig order. This saves scripts the trouble of implementing their own
tree walking code.

Have node_iter() take a 'unique_syms' flag that can be enabled to only
include symbols defined in multiple locations once. This is often what
you want when generating output (and is used by write_config()). Order
is still preserved.

Piggyback a fix to a syntax error test comment. Parsing has been
tightened up now.

3 files changed, 87 insertions, 43 deletions

diff --git a/examples/print_tree.py b/examples/print_tree.py
index 09883d7..5155cf0 100644
--- a/examples/print_tree.py
+++ b/examples/print_tree.py
@@ -2,6 +2,9 @@
 # sometimes implicitly alter the menu structure (see kconfig-language.txt), and
 # that's implemented too.
 #
+# Note: See the Kconfig.node_iter() function as well, which provides a simpler
+# interface for walking the menu tree.
+#
 # Usage:
 #
 #   $ make [ARCH=<arch>] scriptconfig SCRIPT=Kconfiglib/examples/print_tree.py
diff --git a/kconfiglib.py b/kconfiglib.py
index a943f2f..8abec53 100644
--- a/kconfiglib.py
+++ b/kconfiglib.py
@@ -1087,49 +1087,11 @@ class Kconfig(object):
         with self._open(filename, "w") as f:
             f.write(header)
 
-            # Symbol._visited is set to True when we visit a symbol, so that
-            # symbols defined in multiple locations only get one .config entry.
-            # We reset it prior to writing out a new .config. It only needs to
-            # be reset for defined symbols, because undefined symbols will
-            # never be written out (because they do not appear in the menu tree
-            # rooted at Kconfig.top_node).
-            #
-            # The C tools reuse _write_to_conf for this, but we cache
-            # _write_to_conf together with the value and don't invalidate
-            # cached values when writing .config files, so that won't work.
-            #
-            # Note: The usage of _visited here is completely independent from
-            # the usage during dependency loop detection (which runs at the end
-            # of parsing). The attribute is just reused.
-            for sym in self.unique_defined_syms:
-                sym._visited = False
-
-            # The 'top_node' menu node itself doesn't generate any output, so
-            # it's skipped over below
-            node = self.top_node
-            while 1:
-                # Jump to the next node with an iterative tree walk
-                if node.list:
-                    node = node.list
-                elif node.next:
-                    node = node.next
-                else:
-                    while node.parent:
-                        node = node.parent
-                        if node.next:
-                            node = node.next
-                            break
-                    else:
-                        return
-
-                # Write node
-
+            for node in self.node_iter(unique_syms=True):
                 item = node.item
 
                 if isinstance(item, Symbol):
-                    if not item._visited:
-                        item._visited = True
-                        f.write(item.config_string)
+                    f.write(item.config_string)
 
                 elif expr_value(node.dep) and \
                      ((item == MENU and expr_value(node.visibility)) or
@@ -1354,6 +1316,56 @@ class Kconfig(object):
 
                     self.syms[name]._old_val = val
 
+    def node_iter(self, unique_syms=False):
+        """
+        Returns a generator for iterating through all MenuNode's in the Kconfig
+        tree. The iteration is done in Kconfig definition order (the children
+        of a node are visited before the next node is visited).
+
+        The Kconfig.top_node menu node is skipped. It contains an implicit menu
+        that holds the top-level items.
+
+        As an example, the following code will produce a list equal to
+        Kconfig.defined_syms:
+
+          defined_syms = [node.item for node in kconf.node_iter()
+                          if isinstance(node.item, Symbol)]
+
+        unique_syms (default: False):
+          If True, only the first MenuNode will be included for symbols defined
+          in multiple locations.
+
+          Using kconf.node_iter(True) in the example above would give a list
+          equal to unique_defined_syms.
+        """
+        if unique_syms:
+            for sym in self.unique_defined_syms:
+                sym._visited = False
+
+        node = self.top_node
+        while 1:
+            # Jump to the next node with an iterative tree walk
+            if node.list:
+                node = node.list
+            elif node.next:
+                node = node.next
+            else:
+                while node.parent:
+                    node = node.parent
+                    if node.next:
+                        node = node.next
+                        break
+                else:
+                    # No more nodes
+                    return
+
+            if unique_syms and isinstance(node.item, Symbol):
+                if node.item._visited:
+                    continue
+                node.item._visited = True
+
+            yield node
+
     def eval_string(self, s):
         """
         Returns the tristate value of the expression 's', represented as 0, 1,
diff --git a/testsuite.py b/testsuite.py
index eef8389..1e3d89a 100644
--- a/testsuite.py
+++ b/testsuite.py
@@ -464,9 +464,7 @@ def run_selftests():
             fail('expected eval_string("{}") to throw KconfigError, '
                  "didn't".format(expr))
 
-    # The C implementation's parser can be pretty lax about syntax. Kconfiglib
-    # sometimes needs to emulate that. Verify that some bad stuff throws
-    # KconfigError at least.
+    # Verify that some bad stuff throws KconfigError's
     verify_eval_bad("")
     verify_eval_bad("&")
     verify_eval_bad("|")
@@ -1061,6 +1059,37 @@ g
         else:
             fail("'rsource' with missing file did not raise exception")
 
+
+    print("Testing Kconfig.node_iter()")
+
+    # Reuse tests/Klocation. The node_iter(unique_syms=True) case already gets
+    # plenty of testing from write_config() as well.
+
+    c = Kconfig("tests/Klocation", warn=False)
+
+    verify_equal(
+        [node.item.name for node in c.node_iter()
+         if isinstance(node.item, Symbol)],
+        ["SINGLE_DEF", "MULTI_DEF", "HELP_1", "HELP_2", "HELP_3", "MULTI_DEF",
+         "MULTI_DEF", "MENU_HOOK", "COMMENT_HOOK"] + 10*["MULTI_DEF"])
+
+    verify_equal(
+        [node.item.name for node in c.node_iter(True)
+         if isinstance(node.item, Symbol)],
+        ["SINGLE_DEF", "MULTI_DEF", "HELP_1", "HELP_2", "HELP_3", "MENU_HOOK",
+         "COMMENT_HOOK"])
+
+    verify_equal(
+        [node.prompt[0] for node in c.node_iter()
+         if not isinstance(node.item, Symbol)],
+        ["choice", "menu", "comment"])
+
+    verify_equal(
+        [node.prompt[0] for node in c.node_iter(True)
+         if not isinstance(node.item, Symbol)],
+        ["choice", "menu", "comment"])
+
+    # Get rid of custom 'srctree' from Klocation test
     os.environ.pop("srctree", None)