feat: support multiple response types (oxidecomputer#344, oxidecomput…

…er#395)

Add support for multiple response types in the generated client, including handling JSON deserializable types, empty bodies, streams of bytes, and upgraded connections.

This change addresses the following issues:
- Multiple response types not supported (oxidecomputer#344)
- Path to GitHub client (oxidecomputer#395)

progenitor-client/src/progenitor_client.rs

@@ -78,6 +78,18 @@ impl<T: DeserializeOwned> ResponseValue<T> {
             headers,
         })
     }
+
+    /// Transforms the inner data of this `ResponseValue` using a provided function, returning a new `ResponseValue` with the transformed data.
+    pub fn map_inner<U, F>(self, op: F) -> ResponseValue<U>
+    where
+        F: FnOnce(T) -> U,
+    {
+        ResponseValue {
+            inner: op(self.inner), // Apply the operation to the inner data
+            status: self.status,  // Preserve the status
+            headers: self.headers, // Preserve the headers
+        }
+    }
 }
 
 #[cfg(not(target_arch = "wasm32"))]

progenitor-impl/src/httpmock.rs

@@ -16,6 +16,7 @@ use crate::{
     util::{sanitize, Case},
     validate_openapi, Generator, Result,
 };
+use crate::util::generate_multi_type_identifier;
 
 struct MockOp {
     when: TokenStream,
@@ -313,6 +314,18 @@ impl Generator {
                             },
                         )
                     }
+                    crate::method::OperationResponseKind::Multi(types) => {
+                        let arg_type = generate_multi_type_identifier(types, &self.type_space);
+                        (
+                            quote! {
+                                value: #arg_type,
+                            },
+                            quote! {
+                                .header("content-type", "application/json")
+                                .json_body_obj(value)
+                            },
+                        )
+                    }
                     crate::method::OperationResponseKind::None => {
                         Default::default()
                     }

progenitor-impl/src/lib.rs

@@ -357,6 +357,7 @@ impl Generator {
         }?;
 
         let types = self.type_space.to_stream();
+        let multi_types = self.generate_multi_types_stream(&raw_methods, &self.type_space);
 
         // Generate an implementation of a `Self::as_inner` method, if an inner
         // type is defined.
@@ -427,6 +428,7 @@ impl Generator {
                 use std::convert::TryFrom;
 
                 #types
+                #multi_types
             }
 
             #[derive(Clone, Debug)]

progenitor-impl/src/method.rs

@@ -11,11 +11,11 @@ use proc_macro2::TokenStream;
 use quote::{format_ident, quote, ToTokens};
 use typify::{TypeId, TypeSpace};
 
-use crate::{
-    template::PathTemplate,
-    util::{items, parameter_map, sanitize, unique_ident_from, Case},
-    Error, Generator, Result, TagStyle,
-};
+use crate::{Error, Generator, method, Result, TagStyle, template::PathTemplate, util::{
+    Case, generate_multi_type_for_types_stream, generate_multi_type_identifier, items, parameter_map,
+    sanitize,
+    unique_ident_from,
+}};
 use crate::{to_schema::ToSchema, util::ReferenceOrExt};
 
 /// The intermediate representation of an operation that will become a method.
@@ -262,6 +262,7 @@ pub(crate) enum OperationResponseKind {
     None,
     Raw,
     Upgrade,
+    Multi(Vec<Box<OperationResponseKind>>),
 }
 
 impl OperationResponseKind {
@@ -280,6 +281,10 @@ impl OperationResponseKind {
             OperationResponseKind::Upgrade => {
                 quote! { reqwest::Upgraded }
             }
+            OperationResponseKind::Multi(ref types) => {
+                let type_name = generate_multi_type_identifier(types, type_space);
+                quote! { types::#type_name }
+            }
         }
     }
 }
@@ -1032,9 +1037,22 @@ impl Generator {
                             ResponseValue::upgrade(#response_ident).await
                         }
                     }
+                    OperationResponseKind::Multi(_) => {
+                        panic!("Shouldn't occur for the original response")
+                    }
                 };
 
-                quote! { #pat => { #decode } }
+                match &response_type {
+                    OperationResponseKind::Multi(types) => {
+                        let multi_type_name = generate_multi_type_identifier(
+                            &types,
+                            &self.type_space,
+                        );
+                        let type_name = &response.typ.clone().into_tokens(&self.type_space);
+                        quote! { #pat => { #decode.map(|v: ResponseValue<#type_name>| v.map_inner(types::#multi_type_name::from)) } }
+                    }
+                    _ => { quote! { #pat => { #decode } } }
+                }
             });
 
         // Errors...
@@ -1095,9 +1113,22 @@ impl Generator {
                             );
                         }
                     }
+                    OperationResponseKind::Multi(_) => {
+                        panic!("Shouldn't occur for the original response")
+                    }
                 };
 
-                quote! { #pat => { #decode } }
+                match &response_type {
+                    OperationResponseKind::Multi(types) => {
+                        let multi_type_name = generate_multi_type_identifier(
+                            &types,
+                            &self.type_space,
+                        );
+                        let type_name = &response.typ.clone().into_tokens(&self.type_space);
+                        quote! { #pat => { #decode.map(|v: ResponseValue<#type_name>| v.map_inner(types::#multi_type_name::from)) } }
+                    }
+                    _ => { quote! { #pat => { #decode } } }
+                }
             });
 
         let accept_header = matches!(
@@ -1218,6 +1249,29 @@ impl Generator {
         })
     }
 
+    pub(crate) fn generate_multi_types_stream(
+        &self,
+        input_methods: &[method::OperationMethod],
+        type_space: &TypeSpace) -> TokenStream {
+        let mut streams = Vec::new();
+
+        for method in input_methods {
+            let (success_response_items, response_type) = self.extract_responses(
+                method,
+                method::OperationResponseStatus::is_success_or_default,
+            );
+
+            if let OperationResponseKind::Multi(types) = response_type {
+                let multi_stream = generate_multi_type_for_types_stream(&types, type_space);
+                streams.push(multi_stream);
+            }
+        }
+
+        quote! {
+            #(#streams)*
+        }
+    }
+
     /// Extract responses that match criteria specified by the `filter`. The
     /// result is a `Vec<OperationResponse>` that enumerates the cases matching
     /// the filter, and a `TokenStream` that represents the generated type for
@@ -1261,12 +1315,18 @@ impl Generator {
 
         // TODO to deal with multiple response types, we'll need to create an
         // enum type with variants for each of the response types.
-        assert!(response_types.len() <= 1);
-        let response_type = response_types
-            .into_iter()
-            .next()
-            // TODO should this be OperationResponseType::Raw?
-            .unwrap_or(OperationResponseKind::None);
+        // assert!(response_types.len() <= 1);
+        let response_type = if response_types.len() > 1 {
+            OperationResponseKind::Multi(response_types.into_iter().map(Box::new).collect())
+        } else {
+            response_types
+                .into_iter()
+                .next()
+                // TODO should this be OperationResponseType::Raw?
+                .unwrap_or(OperationResponseKind::None)
+        };
+
+
         (response_items, response_type)
     }
 

progenitor-impl/src/util.rs

@@ -6,9 +6,13 @@ use indexmap::IndexMap;
 use openapiv3::{
     Components, Parameter, ReferenceOr, RequestBody, Response, Schema,
 };
+use proc_macro2::TokenStream;
+use quote::{format_ident, quote};
+use typify::TypeSpace;
 use unicode_ident::{is_xid_continue, is_xid_start};
 
-use crate::Result;
+use crate::method::OperationResponseKind;
+use crate::{Result};
 
 pub(crate) trait ReferenceOrExt<T: ComponentLookup> {
     fn item<'a>(&'a self, components: &'a Option<Components>) -> Result<&'a T>;
@@ -142,3 +146,116 @@ pub(crate) fn unique_ident_from(
         name.insert_str(0, "_");
     }
 }
+
+/// Generates a unique identifier by concatenating the type identifiers for a collection of `OperationResponseKind`.
+/// This function is used to dynamically create enum variant names or type combinations based on the types contained within a `Multi` response kind.
+pub(crate) fn generate_multi_type_identifier(types: &Vec<Box<OperationResponseKind>>, type_space: &TypeSpace) -> TokenStream {
+    let identifiers: Vec<TokenStream> = types.iter()
+        .map(|type_kind| {
+            match type_kind.as_ref() {
+                OperationResponseKind::None => {
+                    // Directly return a TokenStream representing 'None' if the type is None.
+                    // This case handles the scenario where the generated tokens would have been ().
+                    quote! { None }
+                }
+                OperationResponseKind::Upgrade => {
+                    // Directly return a TokenStream representing 'Upgrade' if the type is Upgrade.
+                    // This case handles the scenario where the generated tokens would have been reqwest::Upgraded.
+                    quote! { Upgraded }
+                }
+                OperationResponseKind::Type(type_id) => {
+                    // Directly use the Ident returned from TypeSpace, ensuring no invalid string manipulation
+                    let type_name = format_ident!("{}", type_space.get_type(type_id).unwrap().ident().to_string().replace("types :: ", ""));
+                    quote! { #type_name }
+                }
+                _ => {
+                    // Otherwise, generate tokens normally using the `into_tokens` method.
+                    type_kind.clone().into_tokens(type_space)
+                }
+            }
+        })
+        .collect();
+
+    // Convert each TokenStream to string, concatenate them with "Or", and prepend with "types::"
+    let concatenated_type = identifiers.iter()
+        .map(|ts| ts.to_string())
+        .collect::<Vec<_>>()
+        .join("Or");
+
+    // Parse the concatenated string back to a TokenStream to ensure that it can be used in code generation.
+    // This step assumes that the concatenated string is a valid Rust identifier or code.
+    let tokens = concatenated_type.parse::<TokenStream>().unwrap_or_else(|_| quote! { InvalidIdentifier });
+    quote! { #tokens } // Return the new identifier as a TokenStream
+}
+pub(crate) fn generate_multi_type_for_types_stream(types: &Vec<Box<OperationResponseKind>>, type_space: &TypeSpace) -> TokenStream {
+    let enum_name = generate_multi_type_identifier(types, type_space);
+
+    // Generate enum variants and their `From` implementations
+    let variants: Vec<TokenStream> = types.iter().map(|type_kind| {
+        match type_kind.as_ref() {
+            OperationResponseKind::None => {
+                quote! { None }
+            }
+            OperationResponseKind::Upgrade => {
+                quote! { Upgraded(reqwest::Upgraded) }
+            }
+            OperationResponseKind::Type(type_id) => {
+                let type_ident = type_space.get_type(type_id).unwrap().ident().to_string().replace("types :: ", "").parse::<TokenStream>().unwrap();
+                quote! { #type_ident(#type_ident) }
+            }
+            _ => quote! { Unknown },
+        }
+    }).collect();
+
+    let from_impls: Vec<TokenStream> = types.iter().map(|type_kind| {
+        match type_kind.as_ref() {
+            OperationResponseKind::None => {
+                quote! {
+                    impl From<()> for #enum_name {
+                        fn from(_: ()) -> Self {
+                            #enum_name::None
+                        }
+                    }
+                }
+            }
+            OperationResponseKind::Upgrade => {
+                quote! {
+                    impl From<reqwest::Upgraded> for #enum_name {
+                        fn from(value: reqwest::Upgraded) -> Self {
+                            #enum_name::Upgraded(value)
+                        }
+                    }
+                }
+            }
+            OperationResponseKind::Type(type_id) => {
+                let type_ident = type_space.get_type(type_id).unwrap().ident().to_string().replace("types :: ", "").parse::<TokenStream>().unwrap();
+                quote! {
+                    impl From<#type_ident> for #enum_name {
+                        fn from(value: #type_ident) -> Self {
+                            #enum_name::#type_ident(value)
+                        }
+                    }
+                }
+            }
+            _ => {
+                todo!() // Possibility of nested Multi types given openapi spec?
+            }
+        }
+    }).collect();
+
+    let tokens = quote! {
+        // Define the enum
+        #[derive(Debug)]
+        pub enum #enum_name {
+            #(#variants),*
+        }
+
+        // Define the From implementations
+        #(#from_impls)*
+    };
+
+    println!("Tokens: {}", tokens);
+
+    tokens
+}
+