From 1a2dafc1a5cecea8022d0d0eddb1775f9c182b27 Mon Sep 17 00:00:00 2001
From: Graham Esau <gesau@hotmail.co.uk>
Date: Thu, 15 Apr 2021 18:11:28 +0100
Subject: [PATCH] Handle required flattened Option fields

---
 schemars/src/_private.rs               | 12 ++++++++++--
 schemars/tests/expected/validate.json  |  7 ++++++-
 schemars/tests/validate.rs             |  8 ++++++++
 schemars_derive/src/attr/validation.rs | 18 +++++-------------
 schemars_derive/src/schema_exprs.rs    | 10 ++++++++--
 5 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/schemars/src/_private.rs b/schemars/src/_private.rs
index 8b8e2d1..8428370 100644
--- a/schemars/src/_private.rs
+++ b/schemars/src/_private.rs
@@ -4,9 +4,14 @@ use crate::schema::{Metadata, Schema, SchemaObject};
 use crate::JsonSchema;
 
 // Helper for generating schemas for flattened `Option` fields.
-pub fn json_schema_for_flatten<T: ?Sized + JsonSchema>(gen: &mut SchemaGenerator) -> Schema {
+pub fn json_schema_for_flatten<T: ?Sized + JsonSchema>(
+    gen: &mut SchemaGenerator,
+    required: Option<bool>,
+) -> Schema {
     let mut schema = T::_schemars_private_non_optional_json_schema(gen);
-    if T::_schemars_private_is_option() {
+
+    let required = required.unwrap_or_else(|| !T::_schemars_private_is_option());
+    if !required {
         if let Schema::Object(SchemaObject {
             object: Some(ref mut object_validation),
             ..
@@ -15,6 +20,7 @@ pub fn json_schema_for_flatten<T: ?Sized + JsonSchema>(gen: &mut SchemaGenerator
             object_validation.required.clear();
         }
     }
+
     schema
 }
 
@@ -28,11 +34,13 @@ pub fn add_schema_as_property<T: ?Sized + JsonSchema>(
 ) {
     let is_type_option = T::_schemars_private_is_option();
     let required = required.unwrap_or(!is_type_option);
+
     let mut schema = if required && is_type_option {
         T::_schemars_private_non_optional_json_schema(gen)
     } else {
         gen.subschema_for::<T>()
     };
+
     schema = apply_metadata(schema, metadata);
 
     let object = parent.object();
diff --git a/schemars/tests/expected/validate.json b/schemars/tests/expected/validate.json
index 228f249..f5a6fc2 100644
--- a/schemars/tests/expected/validate.json
+++ b/schemars/tests/expected/validate.json
@@ -14,7 +14,8 @@
     "regex_str1",
     "regex_str2",
     "required_option",
-    "tel"
+    "tel",
+    "x"
   ],
   "properties": {
     "min_max": {
@@ -76,6 +77,10 @@
     },
     "required_option": {
       "type": "boolean"
+    },
+    "x": {
+      "type": "integer",
+      "format": "int32"
     }
   }
 }
\ No newline at end of file
diff --git a/schemars/tests/validate.rs b/schemars/tests/validate.rs
index 6d67531..9992029 100644
--- a/schemars/tests/validate.rs
+++ b/schemars/tests/validate.rs
@@ -32,6 +32,14 @@ pub struct Struct {
     map_contains: HashMap<String, ()>,
     #[validate(required)]
     required_option: Option<bool>,
+    #[validate(required)]
+    #[serde(flatten)]
+    required_flattened: Option<Inner>,
+}
+
+#[derive(Debug, JsonSchema)]
+pub struct Inner {
+    x: i32,
 }
 
 #[test]
diff --git a/schemars_derive/src/attr/validation.rs b/schemars_derive/src/attr/validation.rs
index a721ec5..043cccb 100644
--- a/schemars_derive/src/attr/validation.rs
+++ b/schemars_derive/src/attr/validation.rs
@@ -120,17 +120,9 @@ impl ValidationAttrs {
         self
     }
 
-    pub fn validation_statements(&self, field_name: &str) -> TokenStream {
+    pub fn validation_statements(&self, field_name: &str) -> Option<TokenStream> {
         // Assume that the result will be interpolated in a context with the local variable
         // `schema_object` - the SchemaObject for the struct that contains this field.
-        let mut statements = Vec::new();
-
-        // if self.required {
-        //     statements.push(quote! {
-        //         schema_object.object().required.insert(#field_name.to_owned());
-        //     });
-        // }
-
         let mut array_validation = Vec::new();
         let mut number_validation = Vec::new();
         let mut object_validation = Vec::new();
@@ -210,7 +202,7 @@ impl ValidationAttrs {
             || string_validation.is_some()
             || format.is_some()
         {
-            statements.push(quote! {
+            Some(quote! {
                 if let Some(schemars::schema::Schema::Object(prop_schema_object)) = schema_object
                     .object
                     .as_mut()
@@ -222,10 +214,10 @@ impl ValidationAttrs {
                     #string_validation
                     #format
                 }
-            });
+            })
+        } else {
+            None
         }
-
-        statements.into_iter().collect()
     }
 }
 
diff --git a/schemars_derive/src/schema_exprs.rs b/schemars_derive/src/schema_exprs.rs
index ed69873..5c68c52 100644
--- a/schemars_derive/src/schema_exprs.rs
+++ b/schemars_derive/src/schema_exprs.rs
@@ -449,9 +449,15 @@ fn expr_for_struct(
                 type_defs.push(type_def);
             }
 
-            let gen = quote!(gen);
+            let required = if field.validation_attrs.required {
+                quote!(Some(true))
+            } else {
+                quote!(None)
+            };
+
+            let args = quote!(gen, #required);
             quote_spanned! {ty.span()=>
-                .flatten(schemars::_private::json_schema_for_flatten::<#ty>(#gen))
+                .flatten(schemars::_private::json_schema_for_flatten::<#ty>(#args))
             }
         })
         .collect();