Coverage for fastapi_docx/response_generator.py: 94%

1from typing import Any, TypeVar 

2 

3from fastapi.openapi.constants import REF_TEMPLATE 

4from fastapi.routing import APIRoute 

5from pydantic import BaseModel 

6from pydantic.json_schema import GenerateJsonSchema 

7from starlette.exceptions import HTTPException 

8 

9from fastapi_docx.exception_finder import ErrType 

10 

11ErrSchema = TypeVar("ErrSchema", bound=BaseModel) 

12 

13 

14class HTTPExceptionSchema(BaseModel): 

15 detail: str | None = None 

16 

17 

18def get_model_definition(model: type[BaseModel]) -> tuple[str, dict[str, Any]]: 

19 model_name = model.__name__ 

20 schema_generator = GenerateJsonSchema(by_alias=True, ref_template=REF_TEMPLATE) 

21 m_schema = schema_generator.generate( 

22 model.__pydantic_core_schema__, mode="serialization" 

23 ) 

24 if "description" in m_schema: 

25 m_schema["description"] = m_schema["description"].split("\f")[0] 

26 return model_name, m_schema 

27 

28 

29def add_model_to_openapi(api_schema: dict[str, Any], model: type[BaseModel]) -> None: 

30 model_name, m_schema = get_model_definition(model) 

31 if "components" not in api_schema: 

32 api_schema["components"] = {"schemas": {}} 

33 if "schemas" not in api_schema["components"]: 

34 api_schema["components"]["schemas"] = {} 

35 api_schema["components"]["schemas"][model_name] = m_schema 

36 

37 

38def write_response( 

39 api_schema: dict, 

40 route: APIRoute, 

41 exc: HTTPException, 

42 customError: type[ErrType] | None, 

43 customErrSchema: type[ErrSchema] | None, 

44) -> None: 

45 path = getattr(route, "path") 

46 methods = [method.lower() for method in getattr(route, "methods")] 

47 for method in methods: 

48 status_code = str(exc.status_code) 

49 if status_code not in api_schema["paths"][path][method]["responses"]: 

50 if customError and customErrSchema and isinstance(exc, customError): 

51 api_schema["paths"][path][method]["responses"][status_code] = { 

52 "description": exc.__class__.__name__, 

53 "content": { 

54 "application/json": { 

55 "schema": { 

56 "$ref": f"#/components/schemas/{customErrSchema.__name__}" 

57 } 

58 } 

59 }, 

60 } 

61 else: 

62 api_schema["paths"][path][method]["responses"][status_code] = { 

63 "description": exc.detail, 

64 "content": { 

65 "application/json": { 

66 "schema": { 

67 "$ref": "#/components/schemas/HTTPExceptionSchema" 

68 } 

69 } 

70 }, 

71 }