【DRF开发手册】使用 Django Rest Framework 的 @action 定义自定义方法

本文节选自笔者博客： https://www.blog.zeeland.cn/archives/so3f209hfeac

• 💖 作者简介：大家好，我是Zeeland，全栈领域优质创作者。
• 📝 CSDN主页：Zeeland🔥
• 📣 我的博客：Zeeland
• 📚 Github主页: Undertone0809 (Zeeland) (github.com)
• 🎉 支持我：点赞👍+收藏⭐️+留言📝
• 📣 系列专栏：django开发手册🍁
• 💬介绍：The mixture of software dev+Iot+ml+anything🔥

Django系列专栏

• 【django开发手册】Django 中使用自定义用户模型：一个比自带 User 更强的选择
• 【django开发手册】如何使用select_related进行一次连表查询
• 【django开发手册】drf通过添加自定义字段优化DRF序列化器，轻松实现高速API
• 【django开发手册】解决admin添加外键下拉显示外键的问题
• 【Django Rest Framework优化实践】ResponseResult、异常处理方法详解
• 【DRF】深度分析枚举类型在DRF中的序列化问题及解决方案
• 【django开发手册】关于django admin添加表信息的时候外键无法为空的问题解决方案
• 【django开发手册】django admin如何显示外键对应的字段
• 【django开发手册】DRF外键模型查询没有信息？教你实现序列化返回

【DRF开发手册】使用 Django Rest Framework 的 @action 定义自定义方法

前言

如果你正在使用 Django Rest Framework 来构建 RESTful API，那么你一定会经常使用到 ViewSet 这个工具类。ViewSet 提供了一些常见操作的映射，比如 list、create、retrieve、update、destroy 等，能够很方便地实现 CRUD 操作。

不过，如果你需要实现一些比较特殊的操作，那么 ViewSet 的默认方法就可能无法满足你的需求。例如，在一个投票系统中，我们需要提供一个对某个选项进行投票的接口，此时 ViewSet 默认的 list、create、retrieve 等方法就已经无法胜任此任务了。

在这种情况下，我们可以使用 Django Rest Framework 提供的 @action 装饰器来自定义方法。

如何使用 @action 定义自定义方法

使用 @action 定义自定义方法很简单，只需要在 ViewSet 中定义一个方法，并在该方法上加上 @action 装饰器，就可以将该方法转换为一个 API 接口。例如，我们可以在一个投票系统的 ViewSet 中定义一个投票接口：

from rest_framework.decorators import action
from rest_framework.response import Response

class PollViewSet(viewsets.ModelViewSet):
    queryset = Poll.objects.all()
    serializer_class = PollSerializer

    @action(methods=['post'], detail=True)
    def vote(self, request, pk=None):
        poll = self.get_object()
        option_id = request.data.get('option', None)
        if not option_id:
            return Response({'error': 'option is required'})
        try:
            option = poll.options.get(pk=option_id)
        except Option.DoesNotExist:
            return Response({'error': 'invalid option'})
        option.votes += 1
        option.save()
        return Response({'success': True})

该投票接口仅支持 POST 请求，它将会根据传入的 option 参数来对某个选项进行投票。这里需要注意的是，需要通过 self.get_object() 方法来获取当前的 Poll 对象，需要通过 request.data 来获取 POST 请求的参数。

在定义完自定义方法之后，需要将该方法添加到路由中，这可以通过 Django Rest Framework 提供的默认路由机制来实现：

from rest_framework import routers

router = routers.DefaultRouter()
router.register(r'polls', PollViewSet)

urlpatterns = [
    path('', include(router.urls)),
]

这里的 router.register() 方法将会自动根据 PollViewSet 中定义的方法来生成对应的路由，并将其添加到 urlpatterns 中。

对于 @action 定制方法，Django Rest Framework 会自动生成对应的 URL，也就是说，对于上面的投票例子，URL 是：

/polls/<pk>/vote/

其中 <pk> 对应的是具体的 Poll 对象的 primary key。例如，如果你的 Poll 对象的 primary key 是 3，那么可以这样访问：

/polls/3/vote/

在这个例子中，@action(detail=True) 会自动在 URL 中添加 pk 参数（代表 Poll 对象的 primary key），如果没有设置 detail=True，那么 URL 中就不会出现 pk 参数。由此可以看出，@action(detail=True) 与 @action(detail=False) 的区别在于 URL 中是否需要添加 pk 参数。

至此，我们已经可以使用localhost:8000/polls/<id>/vote来访问这个接口了，如果你在其他定制化的时候不需要传入id参数，那么你可以配置detail=False，这样子就可以使用localhost:8000/polls/vote来访问这个接口了。

测试自定义方法

在实现完自定义方法之后，我们需要对其进行测试，以确保其可以正常工作。这可以通过 Django Rest Framework 提供的 APIClient 来实现。

from django.test import TestCase
from django.urls import reverse
from rest_framework import status
from rest_framework.test import APIClient
from polls.models import Poll, Option

class PollAPITestCase(TestCase):

    def setUp(self):
        self.client = APIClient()

    def test_vote(self):
        poll = Poll.objects.create(title='test poll')
        option = Option.objects.create(poll=poll, title='option1')

        url = reverse('polls-vote', args=[poll.id])
        data = {'option': option.id}
        response = self.client.post(url, data, format='json')

        self.assertEqual(response.status_code, status.HTTP_200_OK)

在这个示例测试中，我们首先使用 Poll.objects.create() 方法创建了一个 Poll 对象和一个 Option 对象，接着使用 Django Rest Framework 提供的 reverse() 函数根据 URL 模式的名称反向生成对应的 URL。最后，我们使用 APIClient 的 post() 方法向该 URL 发送 post 请求，并将选项的主键作为数据传递进去。最终，我们断言响应的 HTTP 状态码是否为 200，以此来判断测试是否通过。

reverse() 函数用来根据 URL 模式的名称生成对应的 URL，它接收一个参数，该参数是 URL 模式的名称，会将该名称反向映射为对应的 URL。例如，reverse('polls-vote', args=[poll.id]) 将会生成 URL /polls/<poll.id>/vote/。如果 URL 中含有不止一个占位符，则可以使用类似于 [poll.id, option.id] 的方式将多个占位符传递到 reverse() 函数中，最终生成具有多个参数的 URL。

总结

在本文中，我们介绍了如何使用 Django Rest Framework 的 @action 装饰器来定义自定义方法，以及如何使用 Django Rest Framework 的 APIClient 对自定义方法进行单元测试。通过本文，你应该已经掌握了 @action 装饰器的用法，以及如何进行单元测试。希望这篇文章能够帮助你更好地理解 Django Rest Framework。