cosmetics_shop/apps/account/models.py

# coding=utf-8

from django.db import models
from django.contrib.auth.models import AbstractBaseUser, PermissionsMixin, BaseUserManager
from django.utils import timezone
from rest_framework.utils import model_meta

from utils.exceptions import CustomError


class UserManager(BaseUserManager):
    def create_employee(self, username, password=None, **extra_fields):
        return self.create_user(User.EMPLOYEE, username, password, **extra_fields)

    def create_customer(self, username, password=None, **extra_fields):
        return self.create_user(User.CUSTOMER, username, password, **extra_fields)

    def create_superuser(self, username, password, **extra_fields):
        u = self.create_employee(username, password, **extra_fields)
        u.is_active = True
        u.is_superuser = True
        u.save(using=self._db)
        return u

    def create_user(self, type, username, password=None, **extra_fields):
        if not username:
            raise CustomError(u'请输入用户名！')
        count = User.objects.filter(username=username).count()
        if count > 0:
            raise CustomError(u'该用户名已存在！')
        user = self.model(
            type=type,
            username=username,
            is_superuser=False,
            last_login=timezone.now(),
            **extra_fields
        )

        user.set_password(password)
        user.save(using=self._db)
        return user


class User(AbstractBaseUser, PermissionsMixin):
    EMPLOYEE = 1
    CUSTOMER = 2
    TYPE_CHOICES = (
        (EMPLOYEE, u'职工'),
        (CUSTOMER, u'客户'),
    )

    type = models.PositiveSmallIntegerField(verbose_name=u"类型", choices=TYPE_CHOICES, default=CUSTOMER, editable=False)
    username = models.CharField(verbose_name=u'帐号', max_length=30, unique=True, db_index=True)
    is_active = models.BooleanField(verbose_name=u'激活', default=True)
    date_joined = models.DateTimeField(verbose_name=u'注册时间', default=timezone.now, editable=False)

    objects = UserManager()

    USERNAME_FIELD = 'username'
    REQUIRED_FIELDS = []

    class Meta:
        db_table = "auth_user"
        verbose_name = u'用户'
        unique_together = [
            ('username')
        ]
        index_together = (
            'date_joined',
        )
        ordering = ['-id']
        default_permissions = ()
        permissions = []

    def __unicode__(self):
        return self.username

    def is_customer(self):
        if self.type & self.CUSTOMER:
            return True
        return False

    def is_employee(self):
        if self.type & self.EMPLOYEE:
            return True
        return False

    def change_password(self, new_password, confirm_password, old_password):
        if new_password != confirm_password:
            raise CustomError(u'两次输入的密码不一致， 请检查')
        if not self.check_password(old_password):
            raise CustomError(u'原密码输入错误， 请检查')
        self.set_password(new_password)

    def update_item(self, validated_data):
        def update():
            info = model_meta.get_field_info(self)
            for attr, value in validated_data.items():
                if attr in info.relations and info.relations[attr].to_many:
                    field = getattr(self, attr)
                    field.set(value)
                else:
                    setattr(self, attr, value)

        if not 'username' in validated_data:
            raise CustomError(u'用户名不能为空！')
        count = User.objects.filter(username=validated_data['username']).exclude(id=self.id).count()
        if count > 0:
            raise CustomError(u'该用户名已存在！')

        if not 'password' in validated_data or not validated_data['password']:
            validated_data['password'] = self.password
            update()
        else:
            update()
            self.set_password(validated_data['password'])
        self.save()
        return self