Get UserModel from get_user_model and do not infer the natural key

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
             ))