Moved function for getting correct user model to compat.py

2 files changed, 16 insertions, 15 deletions

diff --git a/rest_framework/authtoken/models.py b/rest_framework/authtoken/models.py
index 3b8bffeb..4da2aa62 100644
--- a/rest_framework/authtoken/models.py
+++ b/rest_framework/authtoken/models.py
@@ -1,29 +1,16 @@
 import uuid
 import hmac
 from hashlib import sha1
+from rest_framework.compat import User
 from django.db import models
-from django import VERSION
 
 
-if VERSION[:2] in ((1, 5,),):
-    from django.conf import settings
-    if hasattr(settings, 'AUTH_USER_MODEL'):
-        user_model = settings.AUTH_USER_MODEL
-    else:
-        from django.contrib.auth.models import User as user_model
-else:
-    try:
-        from django.db.models.auth import User as user_model
-    except ImportError:
-        raise ImportError('User model is not to be found.')
-    
-
 class Token(models.Model):
     """
     The default authorization token model.
     """
     key = models.CharField(max_length=40, primary_key=True)
-    user = models.OneToOneField(user_model, related_name='auth_token')
+    user = models.OneToOneField(User, related_name='auth_token')
     created = models.DateTimeField(auto_now_add=True)
 
     def save(self, *args, **kwargs):
diff --git a/rest_framework/compat.py b/rest_framework/compat.py
index e38e7c33..d5ad2a7a 100644
--- a/rest_framework/compat.py
+++ b/rest_framework/compat.py
@@ -27,6 +27,20 @@ def get_concrete_model(model_cls):
         return model_cls
 
 
+# Django 1.5 add support for custom auth user model
+if django.VERSION >= (1, 5):
+    from django.conf import settings
+    if hasattr(settings, 'AUTH_USER_MODEL'):
+        User = settings.AUTH_USER_MODEL
+    else:
+        from django.contrib.auth.models import User
+else:
+    try:
+        from django.db.models.auth import User
+    except ImportError:
+        raise ImportError('User model is not to be found.')
+
+
 # First implementation of Django class-based views did not include head method
 # in base View class - https://code.djangoproject.com/ticket/15668
 if django.VERSION >= (1, 4):