Adding ability to flatten threads

Author: rth

jwzthreading/jwzthreading.py

@@ -119,6 +119,19 @@ def size(self):
 
  return 1 + sum([child.size for child in self.children])
 
+ def flatten(self):
+ """ Return a flatten version of the thread
+
+ Returns
+ list [Messages]: a list of messages
+ """
+ from itertools import chain
+
+ in_list = [[self.message]] + [child.flatten() for child in self.children]
+
+ return list(chain.from_iterable(in_list))
+
+
 
 
 class Message(object):

jwzthreading/tests/test_jwz.py

@@ -24,6 +24,7 @@ def test_container():
  assert c.children == []
  assert c.parent is None
  assert not c.has_descendant(c2)
+ assert len(c2.flatten()) == 1
 
  # Add a child
  c3 = Container()
@@ -49,11 +50,13 @@ def test_container():
  assert c3.parent == c2
 
 def test_deep_container():
- """Build a 50000-deep list of nested Containers."""
+ """Build a 100-deep list of nested Containers."""
+
+ N = 100
  parent = Container()
  L = [parent]
 
- for _ in range(50000):
+ for _ in range(N):
  child = Container()
  parent.add_child(child)
  L.append(child)
@@ -65,6 +68,9 @@ def test_deep_container():
  # Test a search that fails
  assert not L[0].has_descendant(Container())
 
+ assert L[0].size == N + 1
+ assert len(L[0].flatten()) == N + 1
+
 
 def test_unique():
  assert unique((1, 2, 3, 1, 2, 3)) == [1, 2, 3]