flatten subset

book/src/config/field.md

@@ -110,3 +110,19 @@ struct Message {
     pub inner: InnerMessageB,
 }
 ```
+
+If you wish to only flatten into only a subset of variants, you can define them like so: 
+
+```
+#[superstruct(variants(A, B))]
+struct InnerMessage {
+    pub x: u64,
+    pub y: u64,
+}
+
+#[superstruct(variants(A, B, C))]
+struct Message {
+    #[superstruct(flatten(A,B))]
+    pub inner: InnerMessage,
+}
+```

src/lib.rs

@@ -1,4 +1,5 @@
 use attributes::{IdentList, NestedMetaList};
+use darling::util::Override;
 use darling::FromMeta;
 use from::{
     generate_from_enum_trait_impl_for_ref, generate_from_variant_trait_impl,
@@ -66,7 +67,7 @@ struct StructOpts {
 #[derive(Debug, Default, FromMeta)]
 struct FieldOpts {
     #[darling(default)]
-    flatten: bool,
+    flatten: Option<Override<HashMap<Ident, ()>>>,
     #[darling(default)]
     only: Option<HashMap<Ident, ()>>,
     #[darling(default)]
@@ -75,6 +76,13 @@ struct FieldOpts {
     partial_getter: Option<GetterOpts>,
 }
 
+fn should_skip(flatten: &Override<HashMap<Ident, ()>>, key: &Ident) -> bool {
+    match flatten {
+        Override::Inherit => false,
+        Override::Explicit(map) => !map.is_empty() && !map.contains_key(key),
+    }
+}
+
 /// Getter configuration for a specific field
 #[derive(Debug, Default, FromMeta)]
 struct GetterOpts {
@@ -195,20 +203,29 @@ pub fn superstruct(args: TokenStream, input: TokenStream) -> TokenStream {
             panic!("can't configure `only` and `getter` on the same field");
         } else if field_opts.only.is_none() && field_opts.partial_getter.is_some() {
             panic!("can't set `partial_getter` options on common field");
-        } else if field_opts.flatten && field_opts.only.is_some() {
+        } else if field_opts.flatten.is_some() && field_opts.only.is_some() {
             panic!("can't set `flatten` and `only` on the same field");
-        } else if field_opts.flatten && field_opts.getter.is_some() {
+        } else if field_opts.flatten.is_some() && field_opts.getter.is_some() {
             panic!("can't set `flatten` and `getter` on the same field");
-        } else if field_opts.flatten && field_opts.partial_getter.is_some() {
+        } else if field_opts.flatten.is_some() && field_opts.partial_getter.is_some() {
             panic!("can't set `flatten` and `partial_getter` on the same field");
         }
 
         let only = field_opts.only.map(|only| only.keys().cloned().collect());
         let getter_opts = field_opts.getter.unwrap_or_default();
         let partial_getter_opts = field_opts.partial_getter.unwrap_or_default();
 
-        if field_opts.flatten {
+        if let Some(flatten_opts) = field_opts.flatten {
             for variant in variant_names {
+                if should_skip(&flatten_opts, variant) {
+                    // Remove the field from the field map
+                    let fields = variant_fields
+                        .get_mut(variant)
+                        .expect("invalid variant name");
+                    fields.remove(field_index);
+                    continue;
+                }
+
                 // Update the struct name for this variant.
                 let mut next_variant_field = output_field.clone();
                 match &mut next_variant_field.ty {

tests/flatten.rs

@@ -60,3 +60,36 @@ fn flatten() {
     let mut message_b_ref_mut = MessageRefMut::B(&mut inner_b);
     assert!(message_b_ref_mut.inner_a_mut().is_err());
 }
+
+#[test]
+fn flatten_subset() {
+    #[superstruct(variants(A, B), variant_attributes(derive(Debug, PartialEq, Eq)))]
+    #[derive(Debug, PartialEq, Eq)]
+    struct InnerMessageSubset {
+        pub x: u64,
+        #[superstruct(only(B))]
+        pub y: u64,
+    }
+
+    #[superstruct(variants(A, B, C), variant_attributes(derive(Debug, PartialEq, Eq)))]
+    #[derive(Debug, PartialEq, Eq)]
+    struct MessageSubset {
+        #[superstruct(flatten(A, B))]
+        pub inner: InnerMessageSubset,
+    }
+
+    let message_a = MessageSubset::A(MessageSubsetA {
+        inner: InnerMessageSubsetA { x: 1 },
+    });
+    let message_b = MessageSubset::B(MessageSubsetB {
+        inner: InnerMessageSubsetB { x: 3, y: 4 },
+    });
+    let message_c = MessageSubset::C(MessageSubsetC {});
+    assert_eq!(message_a.inner_a().unwrap().x, 1);
+    assert!(message_a.inner_b().is_err());
+    assert_eq!(message_b.inner_b().unwrap().x, 3);
+    assert_eq!(message_b.inner_b().unwrap().y, 4);
+    assert!(message_b.inner_a().is_err());
+    assert!(message_c.inner_a().is_err());
+    assert!(message_c.inner_b().is_err());
+}