如何在Django自定义用户模型时定制超级用户创建流程？

问题背景 #

我在应用中定义了如下模型，通过一对一关联扩展Django内置User模型来存储用户地址：

class Address(models.Model):
    tower = models.SmallIntegerField()
    floor = models.SmallIntegerField()
    door = models.CharField(max_length=5)


class User(AbstractUser):
    address = models.OneToOneField(Address, on_delete=models.CASCADE)
    REQUIRED_FIELDS = ["email", "address"]

执行python manage.py createsuperuser时，Django会要求输入Address的ID，但此时数据库中还没有任何地址记录，导致报错：

Username: admin
Email address: admin@admin.com
Address (Address.id): 1
Error: address instance with id 1 does not exist.

希望实现创建超级用户时，直接输入地址的tower、floor、door字段，同时在两张表中创建记录，交互如下：

Username: admin
Email address: admin@admin.com
Tower: 1
Floor: 12
Door: A

解决方案 #

需要通过自定义用户管理器和自定义创建超级用户命令来实现需求，具体步骤如下：

1. 自定义UserManager #

重写用户管理器的create_user和create_superuser方法，让其先创建Address实例，再关联到User：

from django.contrib.auth.models import BaseUserManager
from .models import Address

class UserManager(BaseUserManager):
    def create_user(self, username, email, password=None, **extra_fields):
        if not username:
            raise ValueError("必须设置用户名")
        email = self.normalize_email(email)
        
        # 从额外字段中提取地址信息并创建Address实例
        tower = extra_fields.pop("tower")
        floor = extra_fields.pop("floor")
        door = extra_fields.pop("door")
        address = Address.objects.create(tower=tower, floor=floor, door=door)
        
        # 创建并保存用户
        user = self.model(username=username, email=email, address=address, **extra_fields)
        user.set_password(password)
        user.save(using=self._db)
        return user

    def create_superuser(self, username, email, password=None, **extra_fields):
        # 超级用户权限设置
        extra_fields.setdefault("is_staff", True)
        extra_fields.setdefault("is_superuser", True)

        if extra_fields.get("is_staff") is not True:
            raise ValueError("超级用户必须设置is_staff=True")
        if extra_fields.get("is_superuser") is not True:
            raise ValueError("超级用户必须设置is_superuser=True")

        return self.create_user(username, email, password, **extra_fields)

2. 关联自定义管理器到User模型 #

修改User模型，指定使用自定义的UserManager，并移除REQUIRED_FIELDS中的address（不再需要输入ID）：

from django.contrib.auth.models import AbstractUser
from django.db import models
from .managers import UserManager

class User(AbstractUser):
    address = models.OneToOneField(Address, on_delete=models.CASCADE)
    REQUIRED_FIELDS = ["email"]  # 移除address
    objects = UserManager()  # 指定自定义管理器

3. 自定义createsuperuser命令 #

在应用目录下创建management/commands文件夹，新建createsuperuser.py文件，让命令提示输入地址字段：

from django.contrib.auth.management.commands.createsuperuser import Command as BaseCommand

class Command(BaseCommand):
    def handle(self, *args, **options):
        # 提示用户输入地址相关字段
        tower = self.get_input_data("Tower")
        floor = self.get_input_data("Floor")
        door = self.get_input_data("Door")
        
        # 将地址字段加入额外参数，传递给管理器
        options["extra_fields"]["tower"] = int(tower)
        options["extra_fields"]["floor"] = int(floor)
        options["extra_fields"]["door"] = door
        
        # 调用原命令逻辑创建超级用户
        super().handle(*args, **options)

4. 验证效果 #

执行python manage.py createsuperuser，会按照预期提示输入地址字段，完成后同时在Address和User表中创建对应记录。

内容的提问来源于stack exchange，提问作者Dani