深入理解drf-yasg中的自定义Schema生成

深入理解drf-yasg中的自定义Schema生成

项目概述

drf-yasg是一个强大的工具，能够为Django REST framework项目自动生成OpenAPI(Swagger)文档。本文将重点介绍该项目中自定义Schema生成的高级功能，帮助开发者更好地控制API文档的生成过程。

排除特定端点

在实际开发中，我们可能不希望某些API端点出现在Swagger文档中。drf-yasg提供了两种方式来实现这一需求：

1. 类级别排除：通过设置视图类的swagger_schema属性为None，可以排除该视图的所有方法
2. 方法级别排除：使用@swagger_auto_schema装饰器，并设置auto_schema=None来排除特定操作

# 排除整个视图类
class UserList(APIView):
    swagger_schema = None
    ...

# 仅排除PUT方法
@swagger_auto_schema(method='put', auto_schema=None)
@swagger_auto_schema(methods=['get'], ...)
@api_view(['GET', 'PUT'])
def user_detail(request, pk):
    pass

@swagger_auto_schema装饰器详解

@swagger_auto_schema是drf-yasg中最常用的装饰器，它允许我们覆盖生成的Operation对象的属性。根据不同类型的视图，使用方法略有不同：

1. 基于函数的API视图

对于使用@api_view装饰的函数视图，由于同一个视图可以处理多个HTTP方法，因此需要为每个方法单独添加装饰器：

@swagger_auto_schema(method='get', manual_parameters=[test_param], responses={200: user_response})
@swagger_auto_schema(methods=['put', 'post'], request_body=UserSerializer)
@api_view(['GET', 'PUT', 'POST'])
def user_detail(request, pk):
    ...

2. 基于类的API视图

对于APIView、GenericAPIView等类视图，需要装饰每个操作方法：

class UserList(APIView):
    @swagger_auto_schema(responses={200: UserSerializer(many=True)})
    def get(self, request):
        ...

3. ViewSet视图

对于ViewSet、ModelViewSet等视图集，需要装饰动作方法（如list、create等）。对于自定义的@action，可能需要为每个HTTP方法单独装饰：

class ArticleViewSet(viewsets.ModelViewSet):
    @swagger_auto_schema(operation_description='GET /articles/today/')
    @action(detail=False, methods=['get'])
    def today(self, request):
        ...

高级技巧

1. 使用method_decorator

如果不想重写父类方法，可以使用Django的method_decorator：

@method_decorator(name='list', decorator=swagger_auto_schema(
    operation_description="自定义描述"
))
class ArticleViewSet(viewsets.ModelViewSet):
    ...

2. 直接装饰as_view结果

可以像装饰@api_view一样直接装饰视图的as_view()结果：

decorated_login_view = swagger_auto_schema(
    method='post',
    responses={200: LoginResponseSerializer}
)(LoginView.as_view())

SerializerMethodField支持

drf-yasg提供了两种方式来处理SerializerMethodField的Schema生成：

1. 使用swagger_serializer_method装饰器

@swagger_serializer_method(serializer_or_field=OtherStuffSerializer)
def get_other_stuff(self, obj):
    return OtherStuffSerializer().data

2. 使用Python类型提示

对于简单类型，可以直接使用类型提示：

def get_some_number(self, obj) -> float:
    return 1.0

Serializer Meta类配置

可以在Serializer的Meta类中定义一些选项来控制Schema生成：

class WhateverSerializer(Serializer):
    class Meta:
        ref_name = "CustomName"  # 自定义模型定义名称
        swagger_schema_fields = {
            "type": openapi.TYPE_OBJECT,
            "properties": {...}
        }

高级定制

1. 继承SwaggerAutoSchema

通过继承SwaggerAutoSchema可以实现更高级的控制，例如将所有操作ID转换为驼峰命名：

class CamelCaseOperationIDAutoSchema(SwaggerAutoSchema):
    def get_operation_id(self, operation_keys):
        operation_id = super().get_operation_id(operation_keys)
        return camelize(operation_id, uppercase_first_letter=False)

2. 使用Inspector类

通过实现各种Inspector类（如FieldInspector、SerializerInspector等）可以实现对特定字段、序列化器等的定制行为：

class NoSchemaTitleInspector(FieldInspector):
    def process_result(self, result, method_name, obj, **kwargs):
        if isinstance(result, openapi.Schema.OR_REF):
            schema = openapi.resolve_ref(result, self.components)
            schema.pop('title', None)
        return result

注意事项

1. 引用和Schema对象：OpenAPI中的模型定义是通过引用实现的，同一个序列化器类只会生成一个Schema对象
2. 嵌套字段：ModelSerializer生成的嵌套字段（来自ForeignKey）会以内联方式输出，如需引用方式，需要显式设置序列化器类
3. 避免视图或方法特定的FieldInspector处理引用，因为你无法确定哪个视图会首先生成特定序列化器的Schema

通过本文的介绍，开发者可以充分利用drf-yasg提供的各种定制功能，生成更符合项目需求的API文档。