fixed ugly code

Location header is set just, if there is a Location field on the
serializer.

2 files changed, 9 insertions, 12 deletions

diff --git a/rest_framework/mixins.py b/rest_framework/mixins.py
index eddd8f49..832365e5 100644
--- a/rest_framework/mixins.py
+++ b/rest_framework/mixins.py
@@ -20,17 +20,15 @@ class CreateModelMixin(object):
         if serializer.is_valid():
             self.pre_save(serializer.object)
             self.object = serializer.save()
-            headers = self.get_success_headers(serializer)
+            headers = self.get_success_headers(serializer.data)
             return Response(serializer.data, status=status.HTTP_201_CREATED, headers=headers)
         return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
     
-    def get_success_headers(self,serializer):
-        headers = {}
-        for name,field in serializer.fields.iteritems():
-                if isinstance(field,HyperlinkedIdentityField):
-                    headers["Location"] = field.field_to_native(self.object,name)
-                    break
-        return headers
+    def get_success_headers(self,data):
+        if "url" in data:
+            return {'Location':data.get("url")}
+        else:
+            return {}
     
     def pre_save(self, obj):
         pass
diff --git a/rest_framework/tests/hyperlinkedserializers.py b/rest_framework/tests/hyperlinkedserializers.py
index cc5a19c1..5fc935ae 100644
--- a/rest_framework/tests/hyperlinkedserializers.py
+++ b/rest_framework/tests/hyperlinkedserializers.py
@@ -8,17 +8,16 @@ factory = RequestFactory()
 
 
 class BlogPostCommentSerializer(serializers.ModelSerializer):
-    custom_identity_field = serializers.HyperlinkedIdentityField(view_name='blogpostcomment-detail')
+    url = serializers.HyperlinkedIdentityField(view_name='blogpostcomment-detail')
     text = serializers.CharField()
     blog_post_url = serializers.HyperlinkedRelatedField(source='blog_post', view_name='blogpost-detail')
 
     class Meta:
         model = BlogPostComment
-        fields = ('text', 'blog_post_url')
+        fields = ('text', 'blog_post_url', 'url')
 
 
 class PhotoSerializer(serializers.Serializer):
-    """ When Adding a HyperlinkedIdentityField to this serializer, the TestCreateWithForeignKeysAndCustomSlug will fail """
     description = serializers.CharField()
     album_url = serializers.HyperlinkedRelatedField(source='album', view_name='album-detail', queryset=Album.objects.all(), slug_field='title', slug_url_kwarg='title')
 
@@ -222,7 +221,7 @@ class TestCreateWithForeignKeysAndCustomSlug(TestCase):
         request = factory.post('/photos/', data=data)
         response = self.list_create_view(request).render()
         self.assertEqual(response.status_code, status.HTTP_201_CREATED)
-        self.assertNotIn("Location",response,msg="A Serializer without HyperlinkedIdentityField can not produce a valid Location header (for now). Thats why there shouldn'd be one")
+        self.assertNotIn("Location", response, msg="Location should only be included if there is a 'url' field on the serializer")
         self.assertEqual(self.post.photo_set.count(), 1)
         self.assertEqual(self.post.photo_set.all()[0].description, 'A test photo')