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}