Skip to main content

utils/
schema_document.rs

1use schemars::JsonSchema;
2use schemars::generate::{SchemaGenerator, SchemaSettings};
3use serde_json::{Value, json};
4use std::collections::BTreeSet;
5
6pub struct SchemaDocument {
7    outputs: SchemaGenerator,
8    inputs: SchemaGenerator,
9    external: BTreeSet<String>,
10    roots: Vec<String>,
11}
12
13impl Default for SchemaDocument {
14    fn default() -> Self {
15        Self {
16            outputs: SchemaSettings::default().for_serialize().into_generator(),
17            inputs: SchemaSettings::default().for_deserialize().into_generator(),
18            external: BTreeSet::new(),
19            roots: Vec::new(),
20        }
21    }
22}
23
24impl SchemaDocument {
25    pub fn with_external(register: impl Fn(&mut SchemaGenerator)) -> Self {
26        let mut document = Self::default();
27        register(&mut document.outputs);
28        register(&mut document.inputs);
29        document.external =
30            document.outputs.definitions().keys().chain(document.inputs.definitions().keys()).cloned().collect();
31        document
32    }
33
34    pub fn output<T: JsonSchema>(mut self) -> Self {
35        let root = root::<T>(&mut self.outputs);
36        self.roots.push(root);
37        self
38    }
39
40    pub fn input<T: JsonSchema>(mut self) -> Self {
41        let root = root::<T>(&mut self.inputs);
42        self.roots.push(root);
43        self
44    }
45
46    pub fn build(mut self) -> Value {
47        let mut definitions = self.outputs.take_definitions(true);
48        for (name, schema) in self.inputs.take_definitions(true) {
49            if let Some(output) = definitions.get(&name).filter(|_| !self.external.contains(&name)) {
50                assert_eq!(
51                    output, &schema,
52                    "`{name}` is written and read with different schemas, so it cannot be declared once"
53                );
54            }
55            definitions.insert(name, schema);
56        }
57
58        for name in definitions.keys().filter(|name| !self.external.contains(*name)) {
59            let base = name.trim_end_matches(|character: char| character.is_ascii_digit());
60            assert!(
61                base == name || !self.external.contains(base),
62                "`{base}` names both an external type and one of ours, which schemars renamed to `{name}`; \
63                 give ours a distinct #[schemars(rename)]"
64            );
65        }
66
67        json!({ "roots": self.roots, "$defs": definitions, "external": self.external })
68    }
69
70    pub fn print(self) {
71        println!("{}", serde_json::to_string_pretty(&self.build()).expect("schema document serializes to JSON"));
72    }
73}
74
75fn root<T: JsonSchema>(generator: &mut SchemaGenerator) -> String {
76    let schema = generator.subschema_for::<T>();
77    let reference = schema.get("$ref").and_then(|reference| reference.as_str());
78    reference
79        .and_then(|reference| reference.strip_prefix("#/$defs/"))
80        .expect("a root type has its own definition")
81        .to_owned()
82}