From 1806b82ca870f0e497b7ea352c6ab12d243cb2f1 Mon Sep 17 00:00:00 2001
From: Karl Hobley <karlhobley10@gmail.com>
Date: Mon, 23 Mar 2015 14:48:11 +0000
Subject: [PATCH] Made get_willow_image a context manager

---
 wagtail/wagtailimages/models.py            | 51 ++++++++++++----------
 wagtail/wagtailimages/tests/test_models.py |  9 ++--
 2 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/wagtail/wagtailimages/models.py b/wagtail/wagtailimages/models.py
index 3fa6e390a0..c1844f131a 100644
--- a/wagtail/wagtailimages/models.py
+++ b/wagtail/wagtailimages/models.py
@@ -3,6 +3,7 @@ from __future__ import unicode_literals
 import os.path
 import hashlib
 import re
+from contextlib import contextmanager
 
 from six import BytesIO, text_type
 
@@ -84,6 +85,7 @@ class AbstractImage(models.Model, TagSearchable):
     def __str__(self):
         return self.title
 
+    @contextmanager
     def get_willow_image(self):
         try:
             image_file = self.file.file  # triggers a call to self.storage.open, so IOErrors from missing files will be raised at this point
@@ -95,7 +97,9 @@ class AbstractImage(models.Model, TagSearchable):
         image_file.open('rb')
         image_file.seek(0)
 
-        return WillowImage.open(image_file)
+        yield WillowImage.open(image_file)
+
+        image_file.close()
 
     def get_rect(self):
         return Rect(0, 0, self.width, self.height)
@@ -128,27 +132,27 @@ class AbstractImage(models.Model, TagSearchable):
             self.focal_point_height = None
 
     def get_suggested_focal_point(self):
-        willow = self.get_willow_image()
+        with self.get_willow_image() as willow:
+            faces = willow.detect_faces()
 
-        faces = willow.detect_faces()
-        if faces:
-            # Create a bounding box around all faces
-            left = min(face[0] for face in faces)
-            top = min(face[1] for face in faces)
-            right = max(face[2] for face in faces)
-            bottom = max(face[3] for face in faces)
-            focal_point = Rect(left, top, right, bottom)
-        else:
-            features = willow.detect_features()
-            if features:
-                # Create a bounding box around all features
-                left = min(feature[0] for feature in features)
-                top = min(feature[1] for feature in features)
-                right = max(feature[0] for feature in features)
-                bottom = max(feature[1] for feature in features)
+            if faces:
+                # Create a bounding box around all faces
+                left = min(face[0] for face in faces)
+                top = min(face[1] for face in faces)
+                right = max(face[2] for face in faces)
+                bottom = max(face[3] for face in faces)
                 focal_point = Rect(left, top, right, bottom)
             else:
-                return None
+                features = willow.detect_features()
+                if features:
+                    # Create a bounding box around all features
+                    left = min(feature[0] for feature in features)
+                    top = min(feature[1] for feature in features)
+                    right = max(feature[0] for feature in features)
+                    bottom = max(feature[1] for feature in features)
+                    focal_point = Rect(left, top, right, bottom)
+                else:
+                    return None
 
         # Add 20% to width and height and give it a minimum size
         x, y = focal_point.centroid
@@ -299,12 +303,11 @@ class Filter(models.Model):
         return operations
 
     def run(self, image, output):
-        willow = image.get_willow_image()
+        with image.get_willow_image() as willow:
+            for operation in self.operations:
+                operation.run(willow, image)
 
-        for operation in self.operations:
-            operation.run(willow, image)
-
-        willow.save_as_jpeg(output)
+            willow.save_as_jpeg(output)
 
         return output
 
diff --git a/wagtail/wagtailimages/tests/test_models.py b/wagtail/wagtailimages/tests/test_models.py
index fa34344178..9abb3b093e 100644
--- a/wagtail/wagtailimages/tests/test_models.py
+++ b/wagtail/wagtailimages/tests/test_models.py
@@ -219,9 +219,8 @@ class TestGetWillowImage(TestCase):
         )
 
     def test_willow_image_object_returned(self):
-        willow_image = self.image.get_willow_image()
-
-        self.assertIsInstance(willow_image, WillowImage)
+        with self.image.get_willow_image() as willow_image:
+            self.assertIsInstance(willow_image, WillowImage)
 
     def test_with_missing_image(self):
         # Image id=1 in test fixtures has a missing image file
@@ -229,7 +228,9 @@ class TestGetWillowImage(TestCase):
 
         # Attempting to get the Willow image for images without files
         # should raise a SourceImageIOError
-        self.assertRaises(SourceImageIOError, bad_image.get_willow_image)
+        with self.assertRaises(SourceImageIOError):
+            with bad_image.get_willow_image() as willow_image:
+                self.fail() # Shouldn't get here
 
 
 class TestIssue573(TestCase):