From dce17dbb64d350d53bede3f74bcbd8c9e1e98866 Mon Sep 17 00:00:00 2001
From: Jimmy Cuadra <jimmy@jimmycuadra.com>
Date: Sat, 1 Jul 2017 11:29:23 -0700
Subject: [PATCH] Add support for header fields in responses.

---
 src/api/mod.rs      | 16 +++++++++++++---
 src/api/response.rs | 19 +++++++++++++++----
 src/lib.rs          |  2 +-
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/api/mod.rs b/src/api/mod.rs
index a85f7c39..53cacbe3 100644
--- a/src/api/mod.rs
+++ b/src/api/mod.rs
@@ -145,7 +145,7 @@ impl ToTokens for Api {
                     })
             });
 
-            tokens.append(".and_then(|response_body| {");
+            tokens.append(".and_then(move |response_body| {");
 
             tokens
         } else if self.response.has_body_fields() {
@@ -165,7 +165,7 @@ impl ToTokens for Api {
                     })
             });
 
-            tokens.append(".and_then(|response_body| {");
+            tokens.append(".and_then(move |response_body| {");
 
             tokens
         } else {
@@ -175,7 +175,7 @@ impl ToTokens for Api {
                 let future_response = ::futures::future::ok(())
             });
 
-            tokens.append(".and_then(|_| {");
+            tokens.append(".and_then(move |_| {");
 
             tokens
         };
@@ -183,6 +183,14 @@ impl ToTokens for Api {
         let mut closure_end = Tokens::new();
         closure_end.append("});");
 
+        let extract_headers = if self.response.has_header_fields() {
+            quote! {
+                let mut headers = hyper_response.headers().clone();
+            }
+        } else {
+            Tokens::new()
+        };
+
         let response_init_fields = if self.response.has_fields() {
             self.response.init_fields()
         } else {
@@ -239,6 +247,8 @@ impl ToTokens for Api {
                 #[allow(unused_variables)]
                 fn future_from(hyper_response: ::hyper::Response)
                 -> Box<_Future<Item = Self, Error = Self::Error>> {
+                    #extract_headers
+
                     #deserialize_response_body
 
                     let response = Response {
diff --git a/src/api/response.rs b/src/api/response.rs
index d3c7f434..591341b0 100644
--- a/src/api/response.rs
+++ b/src/api/response.rs
@@ -15,6 +15,10 @@ impl Response {
         self.fields.len() != 0
     }
 
+    pub fn has_header_fields(&self) -> bool {
+        self.fields.iter().any(|field| field.is_header())
+    }
+
     pub fn init_fields(&self) -> Tokens {
         let mut tokens = Tokens::new();
 
@@ -31,11 +35,11 @@ impl Response {
                 ResponseField::Header(ref field) => {
                     let field_name = field.ident.as_ref()
                         .expect("expected body field to have a name");
+                    let field_type = &field.ty;
 
                     tokens.append(quote! {
-                        #field_name: hyper_response.headers()
-                            .get_raw(#field_name)
-                            .expect("missing expected request header: {}", #field_name),
+                        #field_name: headers.remove::<#field_type>()
+                            .expect("missing expected request header"),
                     });
                 }
                 ResponseField::NewtypeBody(ref field) => {
@@ -140,7 +144,7 @@ impl ToTokens for Response {
     fn to_tokens(&self, mut tokens: &mut Tokens) {
         tokens.append(quote! {
             /// Data in the response from this API endpoint.
-            #[derive(Debug, Deserialize)]
+            #[derive(Debug)]
             pub struct Response
         });
 
@@ -199,6 +203,13 @@ impl ResponseField {
             _ => false,
         }
     }
+
+    fn is_header(&self) -> bool {
+        match *self {
+            ResponseField::Header(_) => true,
+            _ => false,
+        }
+    }
 }
 
 enum ResponseFieldKind {
diff --git a/src/lib.rs b/src/lib.rs
index f58ae1f6..74738261 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,7 +2,7 @@
 
 #![deny(missing_debug_implementations)]
 #![feature(proc_macro)]
-#![recursion_limit="128"]
+#![recursion_limit="256"]
 
 extern crate proc_macro;
 #[macro_use] extern crate quote;