Skip to content

Instantly share code, notes, and snippets.

@vinyll
Last active June 3, 2023 17:00
Show Gist options
  • Star 7 You must be signed in to star a gist
  • Fork 2 You must be signed in to fork a gist
  • Save vinyll/6103202 to your computer and use it in GitHub Desktop.
Save vinyll/6103202 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