Coverage for src/lexigram/web/docs/enrichment.py: 18%

1"""OpenAPI enrichment utilities. 

2 

3Adds security schemes, tags, pagination, and error schemas to OpenAPI spec. 

4""" 

5 

6from __future__ import annotations 

7 

8from typing import Any 

9 

10 

11class OpenAPIEnricher: 

12 """Enriches OpenAPI specifications with additional details.""" 

13 

14 def __init__(self, spec: dict[str, Any]): 

15 self.spec = spec 

16 self._ensure_components() 

17 

18 def _ensure_components(self) -> None: 

19 """Ensure components section exists.""" 

20 if "components" not in self.spec: 

21 self.spec["components"] = {} 

22 if "schemas" not in self.spec["components"]: 

23 self.spec["components"]["schemas"] = {} 

24 if "securitySchemes" not in self.spec["components"]: 

25 self.spec["components"]["securitySchemes"] = {} 

26 

27 def add_security_scheme( 

28 self, 

29 name: str, 

30 scheme_type: str, 

31 *, 

32 bearer_format: str | None = None, 

33 flows: dict | None = None, 

34 api_key_name: str | None = None, 

35 api_key_in: str | None = None, 

36 ) -> OpenAPIEnricher: 

37 """Add a security scheme to the spec.""" 

38 schemes = self.spec["components"]["securitySchemes"] 

39 

40 if scheme_type == "http": 

41 schemes[name] = { 

42 "type": "http", 

43 "scheme": "bearer", 

44 } 

45 if bearer_format: 

46 schemes[name]["bearerFormat"] = bearer_format 

47 elif scheme_type == "oauth2": 

48 schemes[name] = { 

49 "type": "oauth2", 

50 "flows": flows or {}, 

51 } 

52 elif scheme_type == "apiKey": 

53 schemes[name] = { 

54 "type": "apiKey", 

55 "name": api_key_name or name, 

56 "in": api_key_in or "header", 

57 } 

58 

59 return self 

60 

61 def add_bearer_auth(self, bearer_format: str = "JWT") -> OpenAPIEnricher: 

62 """Add Bearer authentication (JWT, etc.).""" 

63 return self.add_security_scheme( 

64 "BearerAuth", 

65 "http", 

66 bearer_format=bearer_format, 

67 ) 

68 

69 def add_api_key_auth( 

70 self, 

71 name: str = "ApiKeyAuth", 

72 header_name: str = "X-API-Key", 

73 ) -> OpenAPIEnricher: 

74 """Add API key authentication.""" 

75 return self.add_security_scheme( 

76 name, 

77 "apiKey", 

78 api_key_name=header_name, 

79 api_key_in="header", 

80 ) 

81 

82 def apply_security_to_all_paths( 

83 self, 

84 schemes: list[str], 

85 ) -> OpenAPIEnricher: 

86 """Apply security schemes to all paths.""" 

87 security = [{"scheme": s} for s in schemes] 

88 

89 for _path, path_item in self.spec.get("paths", {}).items(): 

90 for method, operation in path_item.items(): 

91 if method in ("get", "post", "put", "patch", "delete"): 

92 if "security" not in operation: 

93 operation["security"] = security 

94 

95 return self 

96 

97 def add_tag( 

98 self, 

99 name: str, 

100 *, 

101 description: str | None = None, 

102 ) -> OpenAPIEnricher: 

103 """Add a tag to the spec.""" 

104 tags = self.spec.setdefault("tags", []) 

105 

106 # Check if tag already exists 

107 for tag in tags: 

108 if tag["name"] == name: 

109 if description: 

110 tag["description"] = description 

111 return self 

112 

113 tag = {"name": name} 

114 if description: 

115 tag["description"] = description 

116 tags.append(tag) 

117 

118 return self 

119 

120 def add_error_schema( 

121 self, 

122 error_type: str, 

123 status_code: int, 

124 ) -> OpenAPIEnricher: 

125 """Add an error response schema.""" 

126 schema = { 

127 "type": "object", 

128 "properties": { 

129 "error": {"type": "string"}, 

130 "message": {"type": "string"}, 

131 }, 

132 } 

133 

134 # Add to common error schemas 

135 error_name = f"Error{error_type}" 

136 self.spec["components"]["schemas"][error_name] = schema 

137 

138 # Add to each path that uses it 

139 for _path, path_item in self.spec.get("paths", {}).items(): 

140 for method, operation in path_item.items(): 

141 if method in ("get", "post", "put", "patch", "delete"): 

142 responses = operation.get("responses", {}) 

143 if str(status_code) not in responses: 

144 responses[str(status_code)] = { 

145 "description": error_type, 

146 "content": { 

147 "application/json": { 

148 "schema": { 

149 "$ref": f"#/components/schemas/{error_name}", 

150 }, 

151 }, 

152 }, 

153 } 

154 

155 return self 

156 

157 def add_pagination_response( 

158 self, 

159 response_name: str, 

160 item_schema: dict[str, Any], 

161 ) -> OpenAPIEnricher: 

162 """Add a paginated response schema.""" 

163 schema = { 

164 "type": "object", 

165 "properties": { 

166 "items": { 

167 "type": "array", 

168 "items": item_schema, 

169 }, 

170 "meta": { 

171 "type": "object", 

172 "properties": { 

173 "total": {"type": "integer"}, 

174 "page": {"type": "integer"}, 

175 "size": {"type": "integer"}, 

176 "pages": {"type": "integer"}, 

177 "has_next": {"type": "boolean"}, 

178 "has_prev": {"type": "boolean"}, 

179 }, 

180 }, 

181 }, 

182 } 

183 

184 self.spec["components"]["schemas"][response_name] = schema 

185 return self 

186 

187 def add_file_upload_parameter( 

188 self, 

189 name: str, 

190 required: bool = False, 

191 ) -> dict[str, Any]: 

192 """Create a file upload parameter schema.""" 

193 return { 

194 "name": name, 

195 "in": "formData", 

196 "description": "File to upload", 

197 "required": required, 

198 "type": "file", 

199 } 

200 

201 

202def enrich_spec( 

203 spec: dict[str, Any], 

204 *, 

205 add_bearer: bool = True, 

206 add_api_key: bool = False, 

207 add_error_responses: bool = True, 

208) -> dict[str, Any]: 

209 """Enrich an OpenAPI spec with common additions.""" 

210 enricher = OpenAPIEnricher(spec) 

211 

212 if add_bearer: 

213 enricher.add_bearer_auth() 

214 

215 if add_api_key: 

216 enricher.add_api_key_auth() 

217 

218 if add_error_responses: 

219 enricher.add_error_schema("Validation", 400) 

220 enricher.add_error_schema("Unauthorized", 401) 

221 enricher.add_error_schema("Forbidden", 403) 

222 enricher.add_error_schema("NotFound", 404) 

223 enricher.add_error_schema("InternalError", 500) 

224 

225 return spec