Skip to content

Instantly share code, notes, and snippets.

@pije76
Forked from vinyll/models.py
Created August 15, 2017 09:51
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save pije76/ec848f9f391f17f1d17806f368ee13c7 to your computer and use it in GitHub Desktop.
Save pije76/ec848f9f391f17f1d17806f368ee13c7 to your computer and use it in GitHub Desktop.
Advanced user inheritance with Django
from django.db import models
from django.utils import timezone
from django.contrib.auth.models import (AbstractBaseUser,
BaseUserManager as DjBaseUserManager)
from model_utils.managers import InheritanceManager
class BaseUserManager(DjBaseUserManager, InheritanceManager):
"""
Manager for all Users types
create_user() and create_superuser() must be overriden as we do not use
unique username but unique email.
"""
def create_user(self, email=None, password=None, **extra_fields):
now = timezone.now()
email = BaseUserManager.normalize_email(email)
u = GenericUser(email=email, is_superuser=False, last_login=now,
**extra_fields)
u.set_password(password)
u.save(using=self._db)
return u
def create_superuser(self, email, password, **extra_fields):
u = self.create_user(email, password, **extra_fields)
u.is_superuser = True
u.save(using=self._db)
return u
class CallableUser(AbstractBaseUser):
"""
The CallableUser class allows to get any type of user by calling
CallableUser.objects.get_subclass(email="my@email.dom") or
CallableUser.objects.filter(email__endswith="@email.dom").select_subclasses()
"""
objects = BaseUserManager()
class AbstractUser(CallableUser):
"""
Here are the fields that are shared among specific User subtypes.
Making it abstract makes 1 email possible in each User subtype.
"""
email = models.EmailField(unique=True)
is_superuser = False
objects = BaseUserManager()
def __unicode__(self):
return self.email
USERNAME_FIELD = 'email'
REQUIRED_FIELD = USERNAME_FIELD
class Meta:
abstract = True
class GenericUser(AbstractUser):
"""
A GenericUser is any type of system user (such as an admin).
This is the one that should be referenced in settings.AUTH_USER_MODEL
"""
is_superuser = models.BooleanField(default=False)
class Professional(AbstractUser):
"""
User subtype with specific fields and properties
"""
company = models.CharField(max_length=50)
class Individual(AbstractUser):
"""
User subtype with specific fields and properties
"""
name = models.CharField(max_length=50)
from django.test import TestCase
from django.db import IntegrityError
from .models import Individual, Professional, GenericUser, CallableUser
class SimpleTest(TestCase):
def test_user_has_email(self):
user = Individual.objects.create(email="me@home.test")
self.assertEquals(user.email, "me@home.test")
def test_user_unique_for_type(self):
Individual.objects.create(email="me@home.test")
user2 = Individual(email="me@home.test")
self.assertRaises(IntegrityError, user2.save)
def test_user_multiple_for_types(self):
Individual.objects.create(email="me@home.test")
try:
Professional.objects.create(email="me@home.test")
except IntegrityError:
self.fail("A user account could not be saved with this email")
def test_create_user(self):
GenericUser.objects.create_user(email="me@home.test",
password="cantfindthis")
user = GenericUser.objects.get(email="me@home.test")
self.assertFalse(user.is_superuser)
def test_create_superuser(self):
GenericUser.objects.create_superuser(email="me@home.test",
password="cantfindthis")
user = GenericUser.objects.get(email="me@home.test")
self.assertTrue(user.is_superuser)
def test_specific_users_are_not_superusers(self):
user = Individual.objects.create(email="me@home.test")
self.assertFalse(user.is_superuser)
user = Professional.objects.create(email="me@home.test")
self.assertFalse(user.is_superuser)
def test_retrieve_a_user_type(self):
Individual.objects.create(email="me@home.test")
user_pk = Individual.objects.get(email="me@home.test").pk
self.assertIsInstance(CallableUser.objects.get_subclass(pk=user_pk),
Individual)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment