Implement JsonSchemaAs for KeyValueMap (#713)

serde_with/src/schemars_0_8.rs

@@ -12,11 +12,15 @@ use crate::{
 use ::schemars_0_8::{
     gen::SchemaGenerator,
     schema::{
-        ArrayValidation, InstanceType, Metadata, NumberValidation, Schema, SchemaObject,
-        SubschemaValidation,
+        ArrayValidation, InstanceType, Metadata, NumberValidation, ObjectValidation, Schema,
+        SchemaObject, SingleOrVec, SubschemaValidation,
     },
     JsonSchema,
 };
+use core::{
+    mem::ManuallyDrop,
+    ops::{Deref, DerefMut},
+};
 
 //===================================================================
 // Trait Definition
@@ -579,6 +583,183 @@ where
     }
 }
 
+impl<T, TA> WrapSchema<Vec<T>, KeyValueMap<TA>>
+where
+    TA: JsonSchemaAs<T>,
+{
+    /// Transform a schema from the entry type of a `KeyValueMap<T>` to the
+    /// resulting field type.
+    ///
+    /// This usually means doing one of two things:
+    /// 1. removing the `$key$` property from an object, or,
+    /// 2. removing the first item from an array.
+    ///
+    /// We also need to adjust any fields that control the number of items or
+    /// properties allowed such as `(max|min)_properties` or `(max|min)_items`.
+    ///
+    /// This is mostly straightforward. Where things get hairy is when dealing
+    /// with subschemas. JSON schemas allow you to build the schema for an
+    /// object by combining multiple subschemas:
+    /// - You can match exactly one of a set of subschemas (`one_of`).
+    /// - You can match any of a set of subschemas (`any_of`).
+    /// - You can match all of a set of subschemas (`all_of`).
+    ///
+    /// Unfortunately for us, we need to handle all of these options by recursing
+    /// into the subschemas and applying the same transformations as above.
+    fn kvmap_transform_schema(gen: &mut SchemaGenerator, schema: &mut Schema) {
+        let mut parents = Vec::new();
+
+        Self::kvmap_transform_schema_impl(gen, schema, &mut parents, 0);
+    }
+
+    fn kvmap_transform_schema_impl(
+        gen: &mut SchemaGenerator,
+        schema: &mut Schema,
+        parents: &mut Vec<String>,
+        depth: u32,
+    ) {
+        if depth > 8 {
+            return;
+        }
+
+        let mut done = false;
+        let schema = match schema {
+            Schema::Object(schema) => schema,
+            _ => return,
+        };
+
+        // The schema is a reference to a schema defined elsewhere.
+        //
+        // If possible we replace it with its definition but if that is not
+        // available then we give up and leave it as-is.
+        let mut parents = if let Some(reference) = &schema.reference {
+            let name = match reference.strip_prefix(&gen.settings().definitions_path) {
+                Some(name) => name,
+                // Reference is defined elsewhere, nothing we can do.
+                None => return,
+            };
+
+            // We are in a recursive reference loop. No point in continuing.
+            if parents.iter().any(|parent| parent == name) {
+                return;
+            }
+
+            let name = name.to_owned();
+            *schema = match gen.definitions().get(&name) {
+                Some(Schema::Object(schema)) => schema.clone(),
+                _ => return,
+            };
+
+            parents.push(name);
+            DropGuard::new(parents, |parents| drop(parents.pop()))
+        } else {
+            DropGuard::unguarded(parents)
+        };
+
+        if let Some(object) = &mut schema.object {
+            // For objects KeyValueMap uses the $key$ property so we need to remove it from
+            // the inner schema.
+
+            done |= object.properties.remove("$key$").is_some();
+            done |= object.required.remove("$key$");
+
+            if let Some(max) = &mut object.max_properties {
+                *max = max.saturating_sub(1);
+            }
+
+            if let Some(min) = &mut object.max_properties {
+                *min = min.saturating_sub(1);
+            }
+        }
+
+        if let Some(array) = &mut schema.array {
+            // For arrays KeyValueMap uses the first array element so we need to remove it
+            // from the inner schema.
+
+            if let Some(SingleOrVec::Vec(items)) = &mut array.items {
+                // If the array is empty then the leading element may be following the
+                // additionalItem schema. In that case we do nothing.
+                if !items.is_empty() {
+                    items.remove(0);
+                    done = true;
+                }
+            }
+
+            if let Some(max) = &mut array.max_items {
+                *max = max.saturating_sub(1);
+            }
+
+            if let Some(min) = &mut array.min_items {
+                *min = min.saturating_sub(1);
+            }
+        }
+
+        // We've already modified the schema so there's no need to do more work.
+        if done {
+            return;
+        }
+
+        let subschemas = match &mut schema.subschemas {
+            Some(subschemas) => subschemas,
+            None => return,
+        };
+
+        if let Some(one_of) = &mut subschemas.one_of {
+            for subschema in one_of {
+                Self::kvmap_transform_schema_impl(gen, subschema, &mut parents, depth + 1);
+            }
+        }
+
+        if let Some(any_of) = &mut subschemas.any_of {
+            for subschema in any_of {
+                Self::kvmap_transform_schema_impl(gen, subschema, &mut parents, depth + 1);
+            }
+        }
+
+        if let Some(all_of) = &mut subschemas.all_of {
+            for subschema in all_of {
+                Self::kvmap_transform_schema_impl(gen, subschema, &mut parents, depth + 1);
+            }
+        }
+    }
+}
+
+impl<T, TA> JsonSchemaAs<Vec<T>> for KeyValueMap<TA>
+where
+    TA: JsonSchemaAs<T>,
+{
+    fn schema_name() -> String {
+        std::format!("KeyValueMap<{}>", <WrapSchema<T, TA>>::schema_name())
+    }
+
+    fn schema_id() -> Cow<'static, str> {
+        std::format!(
+            "serde_with::KeyValueMap<{}>",
+            <WrapSchema<T, TA>>::schema_id()
+        )
+        .into()
+    }
+
+    fn json_schema(gen: &mut SchemaGenerator) -> Schema {
+        let mut value = <WrapSchema<T, TA>>::json_schema(gen);
+        <WrapSchema<Vec<T>, KeyValueMap<TA>>>::kvmap_transform_schema(gen, &mut value);
+
+        SchemaObject {
+            instance_type: Some(InstanceType::Object.into()),
+            object: Some(Box::new(ObjectValidation {
+                additional_properties: Some(Box::new(value)),
+                ..Default::default()
+            })),
+            ..Default::default()
+        }
+        .into()
+    }
+
+    fn is_referenceable() -> bool {
+        true
+    }
+}
+
 impl<K, V, KA, VA, const N: usize> JsonSchemaAs<[(K, V); N]> for Map<KA, VA>
 where
     VA: JsonSchemaAs<V>,
@@ -904,3 +1085,52 @@ forward_duration_schema!(TimestampSecondsWithFrac);
 forward_duration_schema!(TimestampMilliSecondsWithFrac);
 forward_duration_schema!(TimestampMicroSecondsWithFrac);
 forward_duration_schema!(TimestampNanoSecondsWithFrac);
+
+//===================================================================
+// Extra internal helper structs
+
+struct DropGuard<T, F: FnOnce(T)> {
+    value: ManuallyDrop<T>,
+    guard: Option<F>,
+}
+
+impl<T, F: FnOnce(T)> DropGuard<T, F> {
+    pub fn new(value: T, guard: F) -> Self {
+        Self {
+            value: ManuallyDrop::new(value),
+            guard: Some(guard),
+        }
+    }
+
+    pub fn unguarded(value: T) -> Self {
+        Self {
+            value: ManuallyDrop::new(value),
+            guard: None,
+        }
+    }
+}
+
+impl<T, F: FnOnce(T)> Deref for DropGuard<T, F> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.value
+    }
+}
+
+impl<T, F: FnOnce(T)> DerefMut for DropGuard<T, F> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.value
+    }
+}
+
+impl<T, F: FnOnce(T)> Drop for DropGuard<T, F> {
+    fn drop(&mut self) {
+        // SAFETY: value is known to be initialized since we only ever remove it here.
+        let value = unsafe { ManuallyDrop::take(&mut self.value) };
+
+        if let Some(guard) = self.guard.take() {
+            guard(value);
+        }
+    }
+}