笔记 DRF ModelViewSet 自定义返回

from rest_framework.response import Response
from rest_framework.viewsets import ModelViewSet


class APIResponse(Response):
    """重写Response的返回数据格式"""
    def __init__(self, code=200, msg='成功', data=None, status=None, headers=None, **kwargs):
        dic = {'code': code, 'msg': msg}
        if data:
            dic = {'code': code, 'msg': msg, 'data': data}
        dic.update(kwargs)
        super().__init__(data=dic, status=status, headers=headers)


class BaseModelViewSet(ModelViewSet):
    """重写ModelViewSet的各个方法返回数据的格式"""

    def list(self, request, *args, **kwargs):
        response = super().list(request, *args, **kwargs)
        return APIResponse(data=response.data)

    def destroy(self, request, *args, **kwargs):
        response = super().destroy(request, *args, **kwargs)
        return APIResponse(status=response.status_code)

    def update(self, request, *args, **kwargs):
        response = super().update(request, *args, **kwargs)
        return APIResponse(data=response.data)

    def retrieve(self, request, *args, **kwargs):
        response = super().retrieve(request, *args, **kwargs)
        return APIResponse(data=response.data)

    def create(self, request, *args, **kwargs):
        response = super().create(request, *args, **kwargs)
        return APIResponse(data=response.data, status=response.status_code, headers=response.headers)