From 57e8583e2958c888d1d3c1df4bfb871e7956e308 Mon Sep 17 00:00:00 2001
From: Ian Cordasco <graffatcolmingov@gmail.com>
Date: Tue, 29 Dec 2015 08:07:47 -0600
Subject: [PATCH] Add tree/sub-tree traversal to Trie

Accurately retrieve listeners for a specific error code
---
 flake8/_trie.py    | 31 +++++++++++++++++++++++++++----
 flake8/notifier.py |  9 +++++----
 2 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/flake8/_trie.py b/flake8/_trie.py
index ddd3c19..858c067 100644
--- a/flake8/_trie.py
+++ b/flake8/_trie.py
@@ -35,6 +35,16 @@ class Trie(object):
             node = child
         return node
 
+    def traverse(self):
+        """Traverse this tree.
+
+        This performs a depth-first pre-order traversal of children in this
+        tree. It returns the results consistently by first sorting the
+        children based on their prefix and then traversing them in
+        alphabetical order.
+        """
+        return self.root.traverse()
+
 
 class TrieNode(object):
     """The majority of the implementation details of a Trie."""
@@ -47,8 +57,8 @@ class TrieNode(object):
 
     def __repr__(self):
         """Generate an easy to read representation of the node."""
-        return 'TrieNode(prefix={0}, data={1}, children={2})'.format(
-            self.prefix, self.data, dict(self.children)
+        return 'TrieNode(prefix={0}, data={1})'.format(
+            self.prefix, self.data
         )
 
     def find_prefix(self, prefix):
@@ -70,5 +80,18 @@ class TrieNode(object):
         return new_node
 
     def traverse(self):
-        """Traverse children of this node in order."""
-        raise NotImplementedError()
+        """Traverse children of this node.
+
+        This performs a depth-first pre-order traversal of the remaining
+        children in this sub-tree. It returns the results consistently by
+        first sorting the children based on their prefix and then traversing
+        them in alphabetical order.
+        """
+        if not self.children:
+            return
+
+        for prefix in sorted(self.children.keys()):
+            child = self.children[prefix]
+            yield child
+            for child in child.traverse():
+                yield child
diff --git a/flake8/notifier.py b/flake8/notifier.py
index 7b24b2f..625acd7 100644
--- a/flake8/notifier.py
+++ b/flake8/notifier.py
@@ -7,12 +7,13 @@ class Notifier(object):
 
     def listeners_for(self, error_code):
         node = self.listeners.find(error_code)
+        if node is None:
+            return
         for listener in node.data:
             yield listener
-        if node.children:
-            for child in node.traverse():
-                for listener in child.data:
-                    yield listener
+        for child in node.traverse():
+            for listener in child.data:
+                yield listener
 
     def notify(self, error_code, *args, **kwargs):
         for listener in self.listeners_for(error_code):