Add factory.django.ImageField (Closes #52).

7 files changed, 231 insertions, 5 deletions

diff --git a/docs/orms.rst b/docs/orms.rst
index 05166de..611a9ae 100644
--- a/docs/orms.rst
+++ b/docs/orms.rst
@@ -72,6 +72,74 @@ All factories for a Django :class:`~django.db.models.Model` should use the
             [<User: john>, <User: jack>]
 
 
+.. class:: FileField
+
+    Custom declarations for :class:`django.db.models.FileField`
+
+    .. method:: __init__(self, from_path='', from_file='', data=b'', filename='example.dat')
+
+        :param str from_path: Use data from the file located at ``from_path``,
+                              and keep its filename
+        :param file from_file: Use the contents of the provided file object; use its filename
+                               if available
+        :param bytes data: Use the provided bytes as file contents
+        :param str filename: The filename for the FileField
+
+.. note:: If the value ``None`` was passed for the :class:`FileField` field, this will
+          disable field generation:
+
+.. code-block:: python
+
+    class MyFactory(factory.django.DjangoModelFactory):
+        FACTORY_FOR = models.MyModel
+
+        the_file = factory.django.FileField(filename='the_file.dat')
+
+.. code-block:: pycon
+
+    >>> MyFactory(the_file__data=b'uhuh').the_file.read()
+    b'uhuh'
+    >>> MyFactory(the_file=None).the_file
+    None
+
+
+.. class:: ImageField
+
+    Custom declarations for :class:`django.db.models.ImageField`
+
+    .. method:: __init__(self, from_path='', from_file='', filename='example.jpg', width=100, height=100, color='green', format='JPEG')
+
+        :param str from_path: Use data from the file located at ``from_path``,
+                              and keep its filename
+        :param file from_file: Use the contents of the provided file object; use its filename
+                               if available
+        :param str filename: The filename for the ImageField
+        :param int width: The width of the generated image (default: ``100``)
+        :param int height: The height of the generated image (default: ``100``)
+        :param str color: The color of the generated image (default: ``'green'``)
+        :param str format: The image format (as supported by PIL) (default: ``'JPEG'``)
+
+.. note:: If the value ``None`` was passed for the :class:`FileField` field, this will
+          disable field generation:
+
+.. note:: Just as Django's :class:`django.db.models.ImageField` requires the
+          Python Imaging Library, this :class:`ImageField` requires it too.
+
+.. code-block:: python
+
+    class MyFactory(factory.django.DjangoModelFactory):
+        FACTORY_FOR = models.MyModel
+
+        the_image = factory.django.ImageField(color='blue')
+
+.. code-block:: pycon
+
+    >>> MyFactory(the_image__width=42).the_image.width
+    42
+    >>> MyFactory(the_image=None).the_image
+    None
+
+
 Mogo
 ----
 
diff --git a/factory/compat.py b/factory/compat.py
index 7c9b845..f77458f 100644
--- a/factory/compat.py
+++ b/factory/compat.py
@@ -31,10 +31,15 @@ PY2 = (sys.version_info[0] == 2)
 if PY2:  # pragma: no cover
     def is_string(obj):
         return isinstance(obj, (str, unicode))
+
+    from StringIO import StringIO as BytesIO
+
 else:  # pragma: no cover
     def is_string(obj):
         return isinstance(obj, str)
 
+    from io import BytesIO
+
 try:  # pragma: no cover
     # Python >= 3.2
     UTC = datetime.timezone.utc
diff --git a/factory/django.py b/factory/django.py
index 7182be6..04bdbbf 100644
--- a/factory/django.py
+++ b/factory/django.py
@@ -34,8 +34,10 @@ except ImportError as e:  # pragma: no cover
     django_files = None
     import_failure = e
 
+
 from . import base
 from . import declarations
+from .compat import BytesIO
 
 
 class DjangoModelFactory(base.Factory):
@@ -111,14 +113,22 @@ class DjangoModelFactory(base.Factory):
 class FileField(declarations.PostGenerationDeclaration):
     """Helper to fill in django.db.models.FileField from a Factory."""
 
-    def __init__(self, *args, **kwargs):
+    DEFAULT_FILENAME = 'example.dat'
+
+    def __init__(self, **defaults):
         if django_files is None:  # pragma: no cover
             raise import_failure
-        super(FileField, self).__init__(*args, **kwargs)
+        self.defaults = defaults
+        super(FileField, self).__init__()
+
+    def _make_data(self, params):
+        """Create data for the field."""
+        return params.get('data', b'')
 
     def _make_content(self, extraction_context):
         path = ''
-        params = extraction_context.extra
+        params = dict(self.defaults)
+        params.update(extraction_context.extra)
 
         if params.get('from_path') and params.get('from_file'):
             raise ValueError(
@@ -142,13 +152,13 @@ class FileField(declarations.PostGenerationDeclaration):
             path = content.name
 
         else:
-            data = params.get('data', b'')
+            data = self._make_data(params)
             content = django_files.base.ContentFile(data)
 
         if path:
             default_filename = os.path.basename(path)
         else:
-            default_filename = 'example.dat'
+            default_filename = self.DEFAULT_FILENAME
 
         filename = params.get('filename', default_filename)
         return filename, content
@@ -166,3 +176,26 @@ class FileField(declarations.PostGenerationDeclaration):
         finally:
             content.file.close()
         return field_file
+
+
+class ImageField(FileField):
+    DEFAULT_FILENAME = 'example.jpg'
+
+    def _make_data(self, params):
+        # ImageField (both django's and factory_boy's) require PIL.
+        # Try to import it along one of its known installation paths.
+        try:
+            from PIL import Image
+        except ImportError:
+            import Image
+
+        width = params.get('width', 100)
+        height = params.get('height', width)
+        color = params.get('blue')
+        image_format = params.get('format', 'JPEG')
+
+        thumb = Image.new('RGB', (width, height), color)
+        thumb_io = BytesIO()
+        thumb.save(thumb_io, format=image_format)
+        return thumb_io.getvalue()
+
diff --git a/tests/djapp/models.py b/tests/djapp/models.py
index 52acebe..7bb5ace 100644
--- a/tests/djapp/models.py
+++ b/tests/djapp/models.py
@@ -42,3 +42,8 @@ WITHFILE_UPLOAD_DIR = os.path.join(settings.MEDIA_ROOT, WITHFILE_UPLOAD_TO)
 class WithFile(models.Model):
     afile = models.FileField(upload_to=WITHFILE_UPLOAD_TO)
 
+
+
+class WithImage(models.Model):
+    animage = models.ImageField(upload_to=WITHFILE_UPLOAD_TO)
+
diff --git a/tests/test_django.py b/tests/test_django.py
index a74a7ae..4ae2a58 100644
--- a/tests/test_django.py
+++ b/tests/test_django.py
@@ -31,6 +31,16 @@ try:
 except ImportError:  # pragma: no cover
     django = None
 
+try:
+    from PIL import Image
+except ImportError:  # pragma: no cover
+    # Try PIL alternate name
+    try:
+        import Image
+    except ImportError:
+        # OK, not installed
+        Image = None
+
 
 from .compat import is_python2, unittest
 from . import testdata
@@ -99,6 +109,12 @@ class WithFileFactory(factory.django.DjangoModelFactory):
     afile = factory.django.FileField()
 
 
+class WithImageFactory(factory.django.DjangoModelFactory):
+    FACTORY_FOR = models.WithImage
+
+    animage = factory.django.ImageField()
+
+
 @unittest.skipIf(django is None, "Django not installed.")
 class DjangoPkSequenceTestCase(django_test.TestCase):
     def setUp(self):
@@ -261,5 +277,103 @@ class DjangoFileFieldTestCase(unittest.TestCase):
         self.assertFalse(o.afile)
 
 
+@unittest.skipIf(django is None, "Django not installed.")
+@unittest.skipIf(Image is None, "PIL not installed.")
+class DjangoImageFieldTestCase(unittest.TestCase):
+
+    def tearDown(self):
+        super(DjangoImageFieldTestCase, self).tearDown()
+        for path in os.listdir(models.WITHFILE_UPLOAD_DIR):
+            # Remove temporary files written during tests.
+            os.unlink(os.path.join(models.WITHFILE_UPLOAD_DIR, path))
+
+    def test_default_build(self):
+        o = WithImageFactory.build()
+        self.assertIsNone(o.pk)
+        self.assertEqual(100, o.animage.width)
+        self.assertEqual(100, o.animage.height)
+        self.assertEqual('django/example.jpg', o.animage.name)
+
+    def test_default_create(self):
+        o = WithImageFactory.create()
+        self.assertIsNotNone(o.pk)
+        self.assertEqual(100, o.animage.width)
+        self.assertEqual(100, o.animage.height)
+        self.assertEqual('django/example.jpg', o.animage.name)
+
+    def test_with_content(self):
+        o = WithImageFactory.build(animage__width=13, animage__color='blue')
+        self.assertIsNone(o.pk)
+        self.assertEqual(13, o.animage.width)
+        self.assertEqual(13, o.animage.height)
+        self.assertEqual('django/example.jpg', o.animage.name)
+
+    def test_with_file(self):
+        with open(testdata.TESTIMAGE_PATH, 'rb') as f:
+            o = WithImageFactory.build(animage__from_file=f)
+        self.assertIsNone(o.pk)
+        # Image file for a 42x42 green jpeg: 301 bytes long.
+        self.assertEqual(301, len(o.animage.read()))
+        self.assertEqual('django/example.jpeg', o.animage.name)
+
+    def test_with_path(self):
+        o = WithImageFactory.build(animage__from_path=testdata.TESTIMAGE_PATH)
+        self.assertIsNone(o.pk)
+        # Image file for a 42x42 green jpeg: 301 bytes long.
+        self.assertEqual(301, len(o.animage.read()))
+        self.assertEqual('django/example.jpeg', o.animage.name)
+
+    def test_with_file_empty_path(self):
+        with open(testdata.TESTIMAGE_PATH, 'rb') as f:
+            o = WithImageFactory.build(
+                animage__from_file=f,
+                animage__from_path=''
+            )
+        self.assertIsNone(o.pk)
+        # Image file for a 42x42 green jpeg: 301 bytes long.
+        self.assertEqual(301, len(o.animage.read()))
+        self.assertEqual('django/example.jpeg', o.animage.name)
+
+    def test_with_path_empty_file(self):
+        o = WithImageFactory.build(
+            animage__from_path=testdata.TESTIMAGE_PATH,
+            animage__from_file=None,
+        )
+        self.assertIsNone(o.pk)
+        # Image file for a 42x42 green jpeg: 301 bytes long.
+        self.assertEqual(301, len(o.animage.read()))
+        self.assertEqual('django/example.jpeg', o.animage.name)
+
+    def test_error_both_file_and_path(self):
+        self.assertRaises(ValueError, WithImageFactory.build,
+            animage__from_file='fakefile',
+            animage__from_path=testdata.TESTIMAGE_PATH,
+        )
+
+    def test_override_filename_with_path(self):
+        o = WithImageFactory.build(
+            animage__from_path=testdata.TESTIMAGE_PATH,
+            animage__filename='example.foo',
+        )
+        self.assertIsNone(o.pk)
+        # Image file for a 42x42 green jpeg: 301 bytes long.
+        self.assertEqual(301, len(o.animage.read()))
+        self.assertEqual('django/example.foo', o.animage.name)
+
+    def test_existing_file(self):
+        o1 = WithImageFactory.build(animage__from_path=testdata.TESTIMAGE_PATH)
+
+        o2 = WithImageFactory.build(animage=o1.animage)
+        self.assertIsNone(o2.pk)
+        # Image file for a 42x42 green jpeg: 301 bytes long.
+        self.assertEqual(301, len(o2.animage.read()))
+        self.assertEqual('django/example_1.jpeg', o2.animage.name)
+
+    def test_no_file(self):
+        o = WithImageFactory.build(animage=None)
+        self.assertIsNone(o.pk)
+        self.assertFalse(o.animage)
+
+
 if __name__ == '__main__':  # pragma: no cover
     unittest.main()
diff --git a/tests/testdata/__init__.py b/tests/testdata/__init__.py
index 3d1d441..9956610 100644
--- a/tests/testdata/__init__.py
+++ b/tests/testdata/__init__.py
@@ -24,3 +24,4 @@ import os.path
 
 TESTDATA_ROOT = os.path.abspath(os.path.dirname(__file__))
 TESTFILE_PATH = os.path.join(TESTDATA_ROOT, 'example.data')
+TESTIMAGE_PATH = os.path.join(TESTDATA_ROOT, 'example.jpeg')
diff --git a/tests/testdata/example.jpeg b/tests/testdata/example.jpeg
new file mode 100644
index 0000000..28beea9
--- /dev/null
+++ b/tests/testdata/example.jpeg