如何扩展 django 的 createsuperuser 管理命令?

How do I extend django's createsuperuser management command?

我通过添加强制性 Company 字段扩展了我的用户模型,这样每个用户都属于某个公司。 Company 字段是另一个定义公司的模型的外键。

这很好用,除了 ./manage createsuperuser 损坏的事实,因为它不知道添加这个自定义必填字段。

事实证明,我可以将 REQUIRED_FIELDS = ['company'] 添加到我的自定义用户模型中,这将告诉 createsuperuser 请求此字段。这样更好,但是由于公司字段是一个外键,它要求提供公司的 ID,任何人都必须使用 createsuperuser 命令来查找它——这对用户来说不太友好。我希望使用 createsuperuser 命令创建的所有管理员都属于同一家公司,并且如果该公司不存在则自动创建它。

阅读 django 的 documentation on createsuperuser,它提到了我想做的确切事情,以及我想做的确切原因:

You can subclass the management command and override get_input_data() if you want to customize data input and validation. Consult the source code for details on the existing implementation and the method’s parameters. For example, it could be useful if you have a ForeignKey in REQUIRED_FIELDS and want to allow creating an instance instead of entering the primary key of an existing instance.

我只是不确定我应该在项目的哪个位置覆盖 get_input_data() 方法。

编辑:

我扩展用户模型的方法是将其添加到我的一个应用程序 (login_app) 的 models.py 中:

class User(AbstractUser):
    company = models.ForeignKey(Company, on_delete=models.CASCADE, blank=False, null=False)

并将其添加到 settings.py:

AUTH_USER_MODEL = 'login_app.User'

您需要创建一个自定义命令,参见Django的相关文档here。 试试下面的方法:

from django.contrib.auth.management.commands import createsuperuser

class Command(createsuperuser.Command):
    def get_input_data(self, field, message, default=None):
        """
        Override this method if you want to customize data inputs or
        validation exceptions.
        """
        raw_value = input(message)
        if default and raw_value == '':
            raw_value = default
        try:
            val = field.clean(raw_value, None)
        except exceptions.ValidationError as e:
            # here you FK code 

        return val