Add lazy loading to factory.Iterator.

factory.Iterator no longers begins iteration of its argument
on declaration, since this behavior may trigger database query
when that argument is, for instance, a Django queryset.

The ``factory.Iterator``'s argument will only be called when
the containing ``Factory`` is first evaluated; this means that
factories using ``factory.Iterator(models.MyThingy.objects.all())``
will no longer call the database at import time.

6 files changed, 72 insertions, 7 deletions

diff --git a/docs/changelog.rst b/docs/changelog.rst
index 886db0b..0cbd4af 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -11,6 +11,9 @@ ChangeLog
     - Add :attr:`factory.FactoryOptions.rename` to help handle conflicting names (:issue:`206`)
     - Add support for random-yet-realistic values through `fake-factory <https://pypi.python.org/pypi/fake-factory>`_,
       through the :class:`factory.Faker` class.
+    - :class:`factory.Iterator` no longer begins iteration of its argument at import time,
+      thus allowing to pass in a lazy iterator such as a Django queryset
+      (i.e ``factory.Iterator(models.MyThingy.objects.all())``).
 
 .. _v2.5.2:
 
diff --git a/docs/recipes.rst b/docs/recipes.rst
index 70eca46..3cbe6d2 100644
--- a/docs/recipes.rst
+++ b/docs/recipes.rst
@@ -33,6 +33,29 @@ use the :class:`~factory.SubFactory` declaration:
         group = factory.SubFactory(GroupFactory)
 
 
+Choosing from a populated table
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If the target of the :class:`~django.db.models.ForeignKey` should be
+chosen from a pre-populated table
+(e.g :class:`django.contrib.contenttypes.models.ContentType`),
+simply use a :class:`factory.Iterator` on the chosen queryset:
+
+.. code-block:: python
+
+    import factory, factory.django
+    from . import models
+
+    class UserFactory(factory.django.DjangoModelFactory):
+        class Meta:
+            model = models.User
+
+        language = factory.Iterator(models.Language.objects.all())
+
+Here, ``models.Language.objects.all()`` won't be evaluated until the
+first call to ``UserFactory``; thus avoiding DB queries at import time.
+
+
 Reverse dependencies (reverse ForeignKey)
 -----------------------------------------
 
diff --git a/factory/declarations.py b/factory/declarations.py
index 8f2314a..f0dbfe5 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -160,12 +160,19 @@ class Iterator(OrderedDeclaration):
     def __init__(self, iterator, cycle=True, getter=None):
         super(Iterator, self).__init__()
         self.getter = getter
+        self.iterator = None
 
         if cycle:
-            iterator = itertools.cycle(iterator)
-        self.iterator = utils.ResetableIterator(iterator)
+            self.iterator_builder = lambda: utils.ResetableIterator(itertools.cycle(iterator))
+        else:
+            self.iterator_builder = lambda: utils.ResetableIterator(iterator)
 
     def evaluate(self, sequence, obj, create, extra=None, containers=()):
+        # Begin unrolling as late as possible.
+        # This helps with ResetableIterator(MyModel.objects.all())
+        if self.iterator is None:
+            self.iterator = self.iterator_builder()
+
         logger.debug("Iterator: Fetching next value from %r", self.iterator)
         value = next(iter(self.iterator))
         if self.getter is None:
diff --git a/setup.py b/setup.py
index 942fa2c..8ca7e4b 100755
--- a/setup.py
+++ b/setup.py
@@ -44,12 +44,12 @@ setup(
     keywords=['factory_boy', 'factory', 'fixtures'],
     packages=['factory'],
     license='MIT',
-    setup_requires=[
-        'setuptools>=0.8',
-    ],
     install_requires=[
         'fake-factory>=0.5.0',
     ],
+    setup_requires=[
+        'setuptools>=0.8',
+    ],
     tests_require=[
         #'mock',
     ],
@@ -69,7 +69,7 @@ setup(
         "Programming Language :: Python :: 3.4",
         "Programming Language :: Python :: Implementation :: PyPy",
         "Topic :: Software Development :: Testing",
-        "Topic :: Software Development :: Libraries :: Python Modules"
+        "Topic :: Software Development :: Libraries :: Python Modules",
     ],
     test_suite='tests',
     test_loader=test_loader,
diff --git a/tests/test_django.py b/tests/test_django.py
index 2cfb55c..bde8efe 100644
--- a/tests/test_django.py
+++ b/tests/test_django.py
@@ -21,6 +21,7 @@
 """Tests for factory_boy/Django interactions."""
 
 import os
+from .compat import is_python2, unittest, mock
 
 
 try:
@@ -64,7 +65,6 @@ except ImportError:  # pragma: no cover
 import factory
 import factory.django
 
-from .compat import is_python2, unittest, mock
 from . import testdata
 from . import tools
 
diff --git a/tests/test_using.py b/tests/test_using.py
index 6d75531..c7d2b85 100644
--- a/tests/test_using.py
+++ b/tests/test_using.py
@@ -1493,6 +1493,38 @@ class IteratorTestCase(unittest.TestCase):
         for i, obj in enumerate(objs):
             self.assertEqual(i + 10, obj.one)
 
+    def test_iterator_late_loading(self):
+        """Ensure that Iterator doesn't unroll on class creation.
+
+        This allows, for Django objects, to call:
+        foo = factory.Iterator(models.MyThingy.objects.all())
+        """
+        class DBRequest(object):
+            def __init__(self):
+                self.ready = False
+
+            def __iter__(self):
+                if not self.ready:
+                    raise ValueError("Not ready!!")
+                return iter([1, 2, 3])
+
+        # calling __iter__() should crash
+        req1 = DBRequest()
+        with self.assertRaises(ValueError):
+            iter(req1)
+
+        req2 = DBRequest()
+
+        class TestObjectFactory(factory.Factory):
+            class Meta:
+                model = TestObject
+
+            one = factory.Iterator(req2)
+
+        req2.ready = True
+        obj = TestObjectFactory()
+        self.assertEqual(1, obj.one)
+
 
 class BetterFakeModelManager(object):
     def __init__(self, keys, instance):