My models.py
file is as follow:
from django.contrib.auth.models import User
class Shopkeeper(models.Model):
user = models.OneToOneField(User, on_delete=models.CASCADE, primary_key=True)
# ...
class Customer(models.Model):
user = models.OneToOneField(User, on_delete=models.CASCADE, primary_key=True)
# ...
And I have some views which only Customers can access after login, but Shopkeepers cannot. And vice versa.
How can I write decorator for such task?
There is nothing magical about a decorator, it is a function that takes as input the function (or class) to decorate, and makes some changes to it. If we look at the login_required
decorator [GitHub], we see:
def login_required(function=None, redirect_field_name=REDIRECT_FIELD_NAME, login_url=None):
"""
Decorator for views that checks that the user is logged in, redirecting
to the log-in page if necessary.
"""
actual_decorator = user_passes_test(
lambda u: u.is_authenticated,
login_url=login_url,
redirect_field_name=redirect_field_name
)
if function:
return actual_decorator(function)
return actual_decorator
We can thus actually simply make a special case of the user_passes_test
decorator:
from django.contrib.auth.decorators import user_passes_test
def shopkeeper_required(function=None):
def is_shopkeeper(u):
return Shopkeeper.objects.filter(user=u).exists()
actual_decorator = user_passes_test(is_shopkeeper)
if function:
return actual_decorator(function)
else:
return actual_decorator
def customer_required(function=None):
def is_customer(u):
return Customer.objects.filter(user=u).exists()
actual_decorator = user_passes_test(is_customer)
if function:
return actual_decorator(function)
else:
return actual_decorator
You can then for example implement it as:
@login_required
@shopkeeper_required
def some_shopkeeper_view(request):
# ...
pass
@login_required
@customer_required
def some_customer_view(request):
# ...
pass
Note that this @shopkeeper_required
does not really enforces that the user is logged in, although in many cases that will be the case.
EDIT:
We can merge this with a @login_required
(by adding a parameter that acts as a switch to turn this behavior on or off, by default on), like:
from django.contrib.auth.decorators import user_passes_test
def shopkeeper_required(function=None, login_required=True, redirect_field_name=REDIRECT_FIELD_NAME, login_url=None):
def is_shopkeeper(u):
if login_required and not u.is_authenticated:
return False
return Shopkeeper.objects.filter(user=u).exists()
actual_decorator = user_passes_test(
is_shopkeeper,
login_url=login_url,
redirect_field_name=redirect_field_name
)
if function:
return actual_decorator(function)
else:
return actual_decorator
def customer_required(function=None, login_required=True, redirect_field_name=REDIRECT_FIELD_NAME, login_url=None):
def is_customer(u):
if login_required and not u.is_authenticated:
return False
return Customer.objects.filter(user=u).exists()
actual_decorator = user_passes_test(
is_customer,
login_url=login_url,
redirect_field_name=redirect_field_name
)
if function:
return actual_decorator(function)
else:
return actual_decorator