Simplified logic for selections based on REST_USE_KNOX

rest_auth/app_settings.py

@@ -1,10 +1,5 @@
 from django.conf import settings
 
-if getattr(settings, 'REST_USE_KNOX', False):
-    from .utils import create_knox_token as default_create_token
-else:
-    from .utils import default_create_token
-
 from rest_auth.serializers import (
     TokenSerializer as DefaultTokenSerializer,
     KnoxSerializer as DefaultKnoxSerializer,
@@ -14,7 +9,7 @@ from rest_auth.serializers import (
     PasswordResetSerializer as DefaultPasswordResetSerializer,
     PasswordResetConfirmSerializer as DefaultPasswordResetConfirmSerializer,
     PasswordChangeSerializer as DefaultPasswordChangeSerializer)
-from .utils import import_callable
+from .utils import import_callable, default_create_token
 
 create_token = import_callable(
     getattr(settings, 'REST_AUTH_TOKEN_CREATOR', default_create_token))

rest_auth/models.py

@@ -1,12 +1,6 @@
 from django.conf import settings
 
-if getattr(settings, 'REST_USE_KNOX', False):
+from rest_framework.authtoken.models import Token as DefaultTokenModel
-    try:
-        from knox.models import AuthToken as DefaultTokenModel
-    except ImportError:
-        raise ImportError("Install django-rest-knox to use REST_USE_KNOX = True")
-else:
-    from rest_framework.authtoken.models import Token as DefaultTokenModel
 
 from .utils import import_callable
 

rest_auth/registration/views.py

@@ -21,7 +21,7 @@ from rest_auth.app_settings import (TokenSerializer,
 from rest_auth.models import TokenModel
 from rest_auth.registration.serializers import (SocialLoginSerializer,
                                                 VerifyEmailSerializer)
-from rest_auth.utils import jwt_encode
+from rest_auth.utils import create_knox_token, jwt_encode
 from rest_auth.views import LoginView
 from .app_settings import RegisterSerializer
 
@@ -73,6 +73,8 @@ class RegisterView(CreateAPIView):
         user = serializer.save(self.request)
         if getattr(settings, 'REST_USE_JWT', False):
             self.token = jwt_encode(user)  
+        elif getattr(settings, 'REST_USE_KNOX', False):
+            self.token = create_knox_token(user)
         else:
             self.token = create_token(self.token_model, user, serializer)
 

rest_auth/utils.py

@@ -16,8 +16,12 @@ def default_create_token(token_model, user, serializer):
     return token
 
 
-def create_knox_token(token_model, user, serializer):
+def create_knox_token(user):
-    token = token_model.objects.create(user=user)
+    try:
+        from knox.models import AuthToken
+    except ImportError:
+        raise ImportError("django-rest-knox needs to be installed")
+    token = AuthToken.objects.create(user=user)
     return token
 
 

rest_auth/views.py

@@ -22,7 +22,7 @@ from .app_settings import (
     JWTSerializer, create_token
 )
 from .models import TokenModel
-from .utils import jwt_encode
+from .utils import create_knox_token, jwt_encode
 
 if getattr(settings, 'REST_USE_KNOX', False):
     try:
@@ -72,6 +72,8 @@ class LoginView(GenericAPIView):
 
         if getattr(settings, 'REST_USE_JWT', False):
             self.token = jwt_encode(self.user)
+        elif getattr(settings, 'REST_USE_KNOX', False):
+            self.token = create_knox_token(self.user)
         else:
             self.token = create_token(self.token_model, self.user,
                                       self.serializer)