diff --git a/rest_framework/authtoken/management/commands/drf_create_token.py b/rest_framework/authtoken/management/commands/drf_create_token.py index 8419614de..2dc333a6b 100644 --- a/rest_framework/authtoken/management/commands/drf_create_token.py +++ b/rest_framework/authtoken/management/commands/drf_create_token.py @@ -1,13 +1,16 @@ -from django.contrib.auth.models import User +from django.contrib.auth import get_user_model from django.core.management.base import BaseCommand from rest_framework.authtoken.models import Token +UserModel = get_user_model() + + class Command(BaseCommand): help = 'Create DRF Token for a given user' def create_user_token(self, username): - user = User.objects.get(username=username) + user = UserModel._default_manager.get_by_natural_key(username) token = Token.objects.get_or_create(user=user) return token[0] @@ -19,7 +22,7 @@ class Command(BaseCommand): try: token = self.create_user_token(username) - except User.DoesNotExist: + except UserModel.DoesNotExist: print('Cannot create the Token: user {0} does not exist'.format( username ))