Ticket #4948: safe_filename_gen.1.2.diff

django/db/models/base.py

@@ -380 +380 @@
             pass
         filename = field.get_filename(filename)
 
-        # If the filename already exists, keep adding an underscore to the name of
-        # the file until the filename doesn't exist.
-        while os.path.exists(os.path.join(settings.MEDIA_ROOT, filename)):
+        # Generate a unique filename
+        while True:
+            full_filename = os.path.join(settings.MEDIA_ROOT, filename)
             try:
+                # Try and open our filename, but only if it doesn't exist 
+                # and we can open it exclusively. This won't necessarily work on
+                # NFS and some other networked filesystems, but will be ok in normal
+                # cases.
+                fd = os.open(full_filename, os.O_RDWR | os.O_CREAT | os.O_EXCL)
+                break;
+            except OSError, e:
+                # Someone else grabbed this name already
+                pass
+
+            # If the filename already exists, keep adding an underscore to the name of
+            # the file until the filename doesn't exist.
+            try:
                 dot_index = filename.rindex('.')
             except ValueError: # filename has no dot
                 filename += '_'
@@ -393 +406 @@
         # Write the file to disk.
         setattr(self, field.attname, filename)
 
-        full_filename = self._get_FIELD_filename(field)
-        fp = open(full_filename, 'wb')
+        fp = os.fdopen(fd, 'w+b')
         fp.write(raw_contents)
         fp.close()
 

tests/regressiontests/file_field/tests.py

@@ -1 +1 @@
 """
-Tests for file field behavior, and specifically #639, in which Model.save() gets
-called *again* for each FileField. This test will fail if calling an
-auto-manipulator's save() method causes Model.save() to be called more than once.
+Tests for file field behavior.
 """
-
+import tempfile
 import os
+import shutil
 import unittest
-from regressiontests.bug639.models import Photo
+from regressiontests.file_field.models import Photo
 from django.http import QueryDict
 from django.utils.datastructures import MultiValueDict
+from django.conf import settings
 
-class Bug639Test(unittest.TestCase):
-        
+class FileFieldTests(unittest.TestCase):
+    def setUp(self):
+        self.img_content = open(os.path.join(os.path.dirname(__file__), "test.jpg"), "rb").read()
+        self.old_media_root = settings.MEDIA_ROOT
+        settings.MEDIA_ROOT = tempfile.gettempdir()
+    
     def testBug639(self):
         """
         Simulate a file upload and check how many times Model.save() gets called.
+        
+        #639, in which Model.save() gets called *again* for each FileField. This 
+        test will fail if calling an auto-manipulator's save() method causes 
+        Model.save() to be called more than once.
         """
         # Grab an image for testing
-        img = open(os.path.join(os.path.dirname(__file__), "test.jpg"), "rb").read()
         
         # Fake a request query dict with the file
         qd = QueryDict("title=Testing&image=", mutable=True)
         qd["image_file"] = {
             "filename" : "test.jpg",
             "content-type" : "image/jpeg",
-            "content" : img
+            "content" : self.img_content
         }
         
         manip = Photo.AddManipulator()
@@ -34 +41 @@
         # Check the savecount stored on the object (see the model)
         self.assertEqual(p._savecount, 1)
         
+    def testFilenameGeneration_unique(self):
+        """
+        Make sure we're generating unique filenames.
+        Note that upload_to is set to 'file_field_test'
+        """
+        
+        p0 = Photo(title="test0")
+        p0.save_image_file('testfg0.jpg', self.img_content)
+        self.assert_(os.path.exists(p0.get_image_filename()))
+        self.assertEqual(p0.image, os.path.normpath("file_field_test/testfg0.jpg"))
+        # Second file should have an '_' before the extension
+        p1 = Photo(title="test1")
+        p1.save_image_file('testfg0.jpg', self.img_content)
+        self.assert_(os.path.exists(p1.get_image_filename()))
+        self.assertNotEqual(p0.image, p1.image)
+        self.assertEqual(p1.image, os.path.normpath("file_field_test/testfg0_.jpg"))
+        
+        # Now test with no extension
+        p2 = Photo(title="test2")
+        p2.save_image_file('testfg1', self.img_content)
+        self.assert_(os.path.exists(p2.get_image_filename()))
+        self.assertEqual(p2.image, os.path.normpath("file_field_test/testfg1"))
+        # Second file should have an '_'
+        p3 = Photo(title="test1")
+        p3.save_image_file('testfg1', self.img_content)
+        self.assert_(os.path.exists(p3.get_image_filename()))
+        self.assertNotEqual(p2.image, p3.image)
+        self.assertEqual(p3.image, os.path.normpath("file_field_test/testfg1_"))
+        
+    def testFileSaving(self):
+        "Test we actually put the content into the file"
+        p = Photo(title="test")
+        p.save_image_file('tests0.jpg', self.img_content)
+        self.assert_(os.path.exists(p.get_image_filename()))
+        saved_img = open(p.get_image_filename(), 'rb').read()
+        self.assertEqual(self.img_content, saved_img)
+    
     def tearDown(self):
         """
-        Make sure to delete the "uploaded" file to avoid clogging /tmp.
+        Make sure to delete the "uploaded" files to avoid clogging /tmp.
         """
-        p = Photo.objects.get()
-        os.unlink(p.get_image_filename())
- No newline at end of file
+        shutil.rmtree(os.path.join(settings.MEDIA_ROOT, 'file_field_test'))
+        settings.MEDIA_ROOT = self.old_media_root

tests/regressiontests/file_field/models.py

@@ -1 @@
-import tempfile
 from django.db import models
 
 class Photo(models.Model):
     title = models.CharField(max_length=30)
-    image = models.FileField(upload_to=tempfile.gettempdir())
+    image = models.FileField(upload_to='file_field_test')
     
     # Support code for the tests; this keeps track of how many times save() gets
     # called on each instance.