33use std:: collections:: { BTreeMap , HashMap } ;
44use syn:: { Fields , FnArg , Pat , PatType , ReturnType , Type } ;
55use vespera_core:: {
6- route:: { MediaType , Operation , Parameter , ParameterLocation , RequestBody , Response } ,
6+ route:: { Header , MediaType , Operation , Parameter , ParameterLocation , RequestBody , Response } ,
77 schema:: { Reference , Schema , SchemaRef , SchemaType } ,
88} ;
99
@@ -1264,6 +1264,39 @@ fn extract_status_code_tuple(err_ty: &Type) -> Option<(u16, Type)> {
12641264 None
12651265}
12661266
1267+ /// Check whether the provided type is a HeaderMap
1268+ fn is_header_map_type ( ty : & Type ) -> bool {
1269+ if let Type :: Path ( type_path) = ty {
1270+ let path = & type_path. path ;
1271+ if path. segments . is_empty ( ) {
1272+ return false ;
1273+ }
1274+ return path. segments . iter ( ) . any ( |s| s. ident == "HeaderMap" ) ;
1275+ }
1276+ false
1277+ }
1278+
1279+ /// Extract payload type from an Ok tuple and track if headers exist.
1280+ /// The last element of the tuple is always treated as the response body.
1281+ /// Any presence of HeaderMap in the tuple marks headers as present.
1282+ fn extract_ok_payload_and_headers ( ok_ty : & Type ) -> ( Type , Option < HashMap < String , Header > > ) {
1283+ if let Type :: Tuple ( tuple) = ok_ty {
1284+ let payload_ty = tuple. elems . last ( ) . map ( |ty| unwrap_json ( ty) . clone ( ) ) ;
1285+ let has_headers = tuple. elems . iter ( ) . any ( is_header_map_type) ;
1286+
1287+ if let Some ( payload_ty) = payload_ty {
1288+ let headers = if has_headers {
1289+ Some ( HashMap :: new ( ) )
1290+ } else {
1291+ None
1292+ } ;
1293+ return ( payload_ty, headers) ;
1294+ }
1295+ }
1296+
1297+ ( ok_ty. clone ( ) , None )
1298+ }
1299+
12671300/// Analyze return type and convert to Responses map
12681301pub fn parse_return_type (
12691302 return_type : & ReturnType ,
@@ -1288,8 +1321,9 @@ pub fn parse_return_type(
12881321 // Check if it's a Result<T, E>
12891322 if let Some ( ( ok_ty, err_ty) ) = extract_result_types ( ty) {
12901323 // Handle success response (200)
1324+ let ( ok_payload_ty, ok_headers) = extract_ok_payload_and_headers ( & ok_ty) ;
12911325 let ok_schema = parse_type_to_schema_ref_with_schemas (
1292- & ok_ty ,
1326+ & ok_payload_ty ,
12931327 known_schemas,
12941328 struct_definitions,
12951329 ) ;
@@ -1307,7 +1341,7 @@ pub fn parse_return_type(
13071341 "200" . to_string ( ) ,
13081342 Response {
13091343 description : "Successful response" . to_string ( ) ,
1310- headers : None ,
1344+ headers : ok_headers ,
13111345 content : Some ( ok_content) ,
13121346 } ,
13131347 ) ;
@@ -1997,6 +2031,100 @@ mod tests {
19972031 }
19982032 }
19992033
2034+ #[ test]
2035+ fn test_parse_return_type_with_header_map_tuple ( ) {
2036+ let known_schemas = HashMap :: new ( ) ;
2037+ let struct_definitions = HashMap :: new ( ) ;
2038+
2039+ let parsed: syn:: Signature =
2040+ syn:: parse_str ( "fn test() -> Result<(HeaderMap, String), String>" )
2041+ . expect ( "Failed to parse return type" ) ;
2042+
2043+ let responses = parse_return_type ( & parsed. output , & known_schemas, & struct_definitions) ;
2044+
2045+ let ok_response = responses. get ( "200" ) . expect ( "Ok response missing" ) ;
2046+ let ok_content = ok_response
2047+ . content
2048+ . as_ref ( )
2049+ . expect ( "Ok content missing" )
2050+ . get ( "application/json" )
2051+ . expect ( "application/json missing" ) ;
2052+
2053+ if let SchemaRef :: Inline ( schema) = ok_content. schema . as_ref ( ) . unwrap ( ) {
2054+ assert_eq ! ( schema. schema_type, Some ( SchemaType :: String ) ) ;
2055+ } else {
2056+ panic ! ( "Expected inline String schema for Ok type" ) ;
2057+ }
2058+
2059+ assert ! (
2060+ ok_response. headers. is_some( ) ,
2061+ "HeaderMap should set headers"
2062+ ) ;
2063+ }
2064+
2065+ #[ test]
2066+ fn test_parse_return_type_with_status_and_header_map_tuple ( ) {
2067+ let known_schemas = HashMap :: new ( ) ;
2068+ let struct_definitions = HashMap :: new ( ) ;
2069+
2070+ let parsed: syn:: Signature =
2071+ syn:: parse_str ( "fn test() -> Result<(StatusCode, HeaderMap, String), String>" )
2072+ . expect ( "Failed to parse return type" ) ;
2073+
2074+ let responses = parse_return_type ( & parsed. output , & known_schemas, & struct_definitions) ;
2075+
2076+ let ok_response = responses. get ( "200" ) . expect ( "Ok response missing" ) ;
2077+ let ok_content = ok_response
2078+ . content
2079+ . as_ref ( )
2080+ . expect ( "Ok content missing" )
2081+ . get ( "application/json" )
2082+ . expect ( "application/json missing" ) ;
2083+
2084+ if let SchemaRef :: Inline ( schema) = ok_content. schema . as_ref ( ) . unwrap ( ) {
2085+ assert_eq ! ( schema. schema_type, Some ( SchemaType :: String ) ) ;
2086+ } else {
2087+ panic ! ( "Expected inline String schema for Ok type" ) ;
2088+ }
2089+
2090+ assert ! (
2091+ ok_response. headers. is_some( ) ,
2092+ "HeaderMap should set headers"
2093+ ) ;
2094+ }
2095+
2096+ #[ test]
2097+ fn test_parse_return_type_with_mixed_tuple_uses_last_as_body ( ) {
2098+ let known_schemas = HashMap :: new ( ) ;
2099+ let struct_definitions = HashMap :: new ( ) ;
2100+
2101+ // Additional tuple elements before the payload should be ignored; last element is body
2102+ let parsed: syn:: Signature =
2103+ syn:: parse_str ( "fn test() -> Result<(StatusCode, HeaderMap, u32, String), String>" )
2104+ . expect ( "Failed to parse return type" ) ;
2105+
2106+ let responses = parse_return_type ( & parsed. output , & known_schemas, & struct_definitions) ;
2107+
2108+ let ok_response = responses. get ( "200" ) . expect ( "Ok response missing" ) ;
2109+ let ok_content = ok_response
2110+ . content
2111+ . as_ref ( )
2112+ . expect ( "Ok content missing" )
2113+ . get ( "application/json" )
2114+ . expect ( "application/json missing" ) ;
2115+
2116+ if let SchemaRef :: Inline ( schema) = ok_content. schema . as_ref ( ) . unwrap ( ) {
2117+ assert_eq ! ( schema. schema_type, Some ( SchemaType :: String ) ) ;
2118+ } else {
2119+ panic ! ( "Expected inline String schema for Ok type" ) ;
2120+ }
2121+
2122+ assert ! (
2123+ ok_response. headers. is_some( ) ,
2124+ "HeaderMap should set headers"
2125+ ) ;
2126+ }
2127+
20002128 #[ test]
20012129 fn test_parse_return_type_primitive_types ( ) {
20022130 let known_schemas = HashMap :: new ( ) ;
0 commit comments