From 6d4f82b0c737fbe46a6ffe8a11ef56314b05b516 Mon Sep 17 00:00:00 2001
From: Max Khardin <mkhardin@yandex-team.ru>
Date: Mon, 2 Mar 2015 22:56:37 +0300
Subject: [PATCH 1/5] Check that given thing is actually iterable before trying
 to iterate over it

---
 builders/builder.py | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/builders/builder.py b/builders/builder.py
index 4144b71..aea7bfd 100644
--- a/builders/builder.py
+++ b/builders/builder.py
@@ -13,15 +13,18 @@ def flatten(l):
         Generator that flattens iterable infinitely. If an item is iterable, ``flatten`` descends on it.
         If it is callable, it descends on the call result (with no arguments), and it yields the item itself otherwise.
     """
-    for el in l:
-        if isinstance(el, collections.Iterable) and not isinstance(el, basestring):
-            for sub in flatten(el):
-                yield sub
-        elif callable(el):
-            for sub in flatten(el()):
-                yield sub
-        else:
-            yield el
+    if not isinstance(l, collections.Iterable):
+        yield l
+    else:
+        for el in l:
+            if isinstance(el, collections.Iterable) and not isinstance(el, basestring):
+                for sub in flatten(el):
+                    yield sub
+            elif callable(el):
+                for sub in flatten(el()):
+                    yield sub
+            else:
+                yield el
 
 
 class Builder:

From f83b48c716455021275942f871e42fab6ec68c49 Mon Sep 17 00:00:00 2001
From: Max Khardin <mkhardin@yandex-team.ru>
Date: Tue, 3 Mar 2015 15:47:13 +0300
Subject: [PATCH 2/5] Fix flatten logic

---
 builders/builder.py | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/builders/builder.py b/builders/builder.py
index aea7bfd..1d9b8e0 100644
--- a/builders/builder.py
+++ b/builders/builder.py
@@ -13,18 +13,15 @@ def flatten(l):
         Generator that flattens iterable infinitely. If an item is iterable, ``flatten`` descends on it.
         If it is callable, it descends on the call result (with no arguments), and it yields the item itself otherwise.
     """
-    if not isinstance(l, collections.Iterable):
-        yield l
-    else:
+    if isinstance(l, collections.Iterable) and not isinstance(l, basestring):
         for el in l:
-            if isinstance(el, collections.Iterable) and not isinstance(el, basestring):
-                for sub in flatten(el):
-                    yield sub
-            elif callable(el):
-                for sub in flatten(el()):
-                    yield sub
-            else:
-                yield el
+            for sub in flatten(el):
+                yield sub
+    elif callable(l):
+        for sub in flatten(l()):
+            yield sub
+    else:
+        yield l
 
 
 class Builder:

From 108aa8796ffd36105dc810e15098cda84e1dc600 Mon Sep 17 00:00:00 2001
From: Max Khardin <mkhardin@yandex-team.ru>
Date: Tue, 3 Mar 2015 15:47:34 +0300
Subject: [PATCH 3/5] Add tests for new flatten logic

---
 builders/tests/test_builder.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/builders/tests/test_builder.py b/builders/tests/test_builder.py
index a9430f8..2f33eda 100644
--- a/builders/tests/test_builder.py
+++ b/builders/tests/test_builder.py
@@ -1,3 +1,5 @@
+import pytest
+
 from builders.builder import Builder, flatten
 from builders.construct import Unique, Collection, Reused
 from builders.modifiers import InstanceModifier
@@ -175,6 +177,25 @@ def test_flatten():
     assert list(flatten(l)) == [1, 2, 3, 4, 5, 'ololo']
 
 
+def test_flatten_noniterable():
+    y = 100
+    f = flatten(y)
+
+    assert f.next() == y
+    with pytest.raises(StopIteration):
+        f.next()
+
+
+def test_flatten_function_returning_noniterable():
+    def y():
+        return 100
+    f = flatten(y)
+
+    assert f.next() == 100
+    with pytest.raises(StopIteration):
+        f.next()
+
+
 def test_flatten_callable():
     def x():
         return [1, 2, 3]

From e4fca88d54c25c451d00464391b7e173800e28b5 Mon Sep 17 00:00:00 2001
From: Max Khardin <mkhardin@yandex-team.ru>
Date: Tue, 3 Mar 2015 16:34:35 +0300
Subject: [PATCH 4/5] Refactor tests a tiny bit

---
 builders/tests/test_builder.py | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/builders/tests/test_builder.py b/builders/tests/test_builder.py
index 2f33eda..72ed1a1 100644
--- a/builders/tests/test_builder.py
+++ b/builders/tests/test_builder.py
@@ -179,21 +179,14 @@ def test_flatten():
 
 def test_flatten_noniterable():
     y = 100
-    f = flatten(y)
-
-    assert f.next() == y
-    with pytest.raises(StopIteration):
-        f.next()
+    assert list(flatten(y)) == [y]
 
 
 def test_flatten_function_returning_noniterable():
     def y():
         return 100
-    f = flatten(y)
 
-    assert f.next() == 100
-    with pytest.raises(StopIteration):
-        f.next()
+    assert list(flatten(y)) == [100]
 
 
 def test_flatten_callable():

From 2f66e5c84142a8d99d63b22e4cffd03f632b8c8c Mon Sep 17 00:00:00 2001
From: Max Khardin <mkhardin@yandex-team.ru>
Date: Tue, 3 Mar 2015 16:41:46 +0300
Subject: [PATCH 5/5] Remove unused import

---
 builders/tests/test_builder.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/builders/tests/test_builder.py b/builders/tests/test_builder.py
index 72ed1a1..264141f 100644
--- a/builders/tests/test_builder.py
+++ b/builders/tests/test_builder.py
@@ -1,5 +1,3 @@
-import pytest
-
 from builders.builder import Builder, flatten
 from builders.construct import Unique, Collection, Reused
 from builders.modifiers import InstanceModifier