如何在Django中编写自己的装饰器?

Abi*_*ood 1 python django django-views python-decorators

我的models.py文件如下:

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)
# ...
Run Code Online (Sandbox Code Playgroud)

我有一些视图,只有客户登录后才能访问,而店主则不能。反之亦然。如何为此类任务编写装饰器?

Wil*_*sem 5

装饰器没有什么神奇的东西,它是一个将要装饰的函数(或类)作为输入并对其进行一些更改的函数。如果查看login_required装饰器[GitHub]则会看到:

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
Run Code Online (Sandbox Code Playgroud)

因此,我们实际上可以简单地对user_passes_test装饰器进行特殊处理:

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
Run Code Online (Sandbox Code Playgroud)

然后,您可以例如将其实现为:

@login_required
@shopkeeper_required
def some_shopkeeper_view(request):
    # ...
    pass

@login_required
@customer_required
def some_customer_view(request):
    # ...
    pass
Run Code Online (Sandbox Code Playgroud)

请注意,尽管在许多情况下确实如此,但这@shopkeeper_required不能真正强制用户登录。

编辑

我们可以将其与合并@login_required(通过添加充当开关以默认打开或关闭此行为的参数),例如:

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
Run Code Online (Sandbox Code Playgroud)