diff --git a/rest_auth/app_settings.py b/rest_auth/app_settings.py index e0ab429..3843665 100644 --- a/rest_auth/app_settings.py +++ b/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)) diff --git a/rest_auth/models.py b/rest_auth/models.py index 13d4f74..a132f9c 100644 --- a/rest_auth/models.py +++ b/rest_auth/models.py @@ -1,12 +1,6 @@ from django.conf import settings -if getattr(settings, 'REST_USE_KNOX', False): - 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 rest_framework.authtoken.models import Token as DefaultTokenModel from .utils import import_callable diff --git a/rest_auth/registration/views.py b/rest_auth/registration/views.py index 5d7c77d..4f3beee 100644 --- a/rest_auth/registration/views.py +++ b/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) diff --git a/rest_auth/utils.py b/rest_auth/utils.py index dbd7ee5..bf147aa 100644 --- a/rest_auth/utils.py +++ b/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): - token = token_model.objects.create(user=user) +def create_knox_token(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 diff --git a/rest_auth/views.py b/rest_auth/views.py index ffb2f67..3544edd 100644 --- a/rest_auth/views.py +++ b/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)