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

2 files changed, 43 insertions, 5 deletions

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()
+