prutoipa_build/
lib.rs

1mod descriptor;
2mod error;
3mod generator;
4mod package;
5mod package_set;
6mod syntax;
7
8use std::{
9    io::{BufWriter, Write},
10    path::PathBuf,
11};
12
13use descriptor::Descriptor;
14use error::PrutoipaBuildError;
15use generator::{enumeration::generate_enum, message::generate_message};
16use package_set::PackageSet;
17use prost_types::FileDescriptorSet;
18
19#[derive(Debug, Default)]
20pub struct Builder {
21    out_dir: Option<PathBuf>,
22    package_set: PackageSet,
23    generate_enum_values: bool,
24}
25
26impl Builder {
27    /// Create a new `Builder`
28    pub fn new() -> Self {
29        Self::default()
30    }
31
32    /// Configures the output directory where generated Rust files will be written.
33    ///
34    /// If unset, defaults to the `OUT_DIR` environment variable. `OUT_DIR` is set by Cargo when
35    /// executing build scripts, so `out_dir` typically does not need to be configured.
36    pub fn out_dir<P>(&mut self, path: P) -> &mut Self
37    where
38        P: Into<PathBuf>,
39    {
40        self.out_dir = Some(path.into());
41        self
42    }
43
44    /// Generate utoipa enum_values property at enums
45    pub fn generate_enum_values(&mut self) -> &mut Self {
46        self.generate_enum_values = true;
47        self
48    }
49
50    /// Register an encoded `FileDescriptorSet` with this `Builder`
51    pub fn register_descriptors_encoded(
52        &mut self,
53        fds_encoded: &[u8],
54    ) -> Result<&mut Self, error::PrutoipaBuildError> {
55        self.package_set
56            .register_file_descriptor_set_encoded(fds_encoded)?;
57
58        Ok(self)
59    }
60
61    /// Register a `FileDescriptorSet` with this `Builder`
62    pub fn register_descriptors(
63        &mut self,
64        fds: FileDescriptorSet,
65    ) -> Result<&mut Self, error::PrutoipaBuildError> {
66        self.package_set.register_file_descriptor_set(fds)?;
67
68        Ok(self)
69    }
70
71    fn get_out_dir(&self) -> Result<PathBuf, PrutoipaBuildError> {
72        if let Some(out_dir) = self.out_dir.clone() {
73            Ok(out_dir)
74        } else {
75            Ok(std::env::var_os("OUT_DIR")
76                .ok_or(PrutoipaBuildError::OutputDirNotSet)?
77                .into())
78        }
79    }
80
81    pub fn build(&mut self) -> Result<(), PrutoipaBuildError> {
82        let mut output = self.get_out_dir()?;
83        output.push("DUMMY_FILENAME");
84
85        let write_factory = move |package_name: String| {
86            output.set_file_name(format!("{}.utoipa.rs", package_name));
87
88            let file = std::fs::OpenOptions::new()
89                .write(true)
90                .truncate(true)
91                .create(true)
92                .open(&output)?;
93
94            Ok(BufWriter::new(file))
95        };
96
97        let writers = self.generate(write_factory)?;
98        for (_, mut writer) in writers {
99            writer.flush()?;
100        }
101
102        Ok(())
103    }
104
105    pub fn generate<W: Write, F: FnMut(String) -> std::io::Result<W>>(
106        &self,
107        mut write_factory: F,
108    ) -> Result<Vec<(String, W)>, PrutoipaBuildError> {
109        self.package_set
110            .get_packages()
111            .into_iter()
112            .map(|(package_name, mut package)| {
113                let mut writer = write_factory(package_name.clone())?;
114
115                package
116                    .get_descriptors()
117                    .into_iter()
118                    .map(|(descriptor_name, descriptor)| match descriptor {
119                        Descriptor::Message(message) => generate_message(
120                            &mut writer,
121                            package.get_name(),
122                            descriptor_name,
123                            message,
124                        ),
125                        Descriptor::Enum(enum_descriptor) => generate_enum(
126                            &mut writer,
127                            package.get_name(),
128                            descriptor_name,
129                            enum_descriptor,
130                            self.generate_enum_values,
131                        ),
132                    })
133                    .collect::<Result<Vec<()>, PrutoipaBuildError>>()?;
134
135                Ok((package_name, writer))
136            })
137            .collect::<Result<Vec<(String, W)>, PrutoipaBuildError>>()
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    // use super::*;
144    use prost::Message;
145    use prost_types::{
146        field_descriptor_proto::{Label, Type},
147        DescriptorProto, EnumDescriptorProto, EnumValueDescriptorProto, FieldDescriptorProto,
148        FileDescriptorProto, FileDescriptorSet,
149    };
150
151    pub fn get_file_descriptor_proto() -> FileDescriptorProto {
152        FileDescriptorProto {
153            syntax: Some("proto3".to_string()),
154            package: Some("people".to_string()),
155            name: Some("person.proto".to_string()),
156            enum_type: vec![EnumDescriptorProto {
157                name: Some("GENDER".to_string()),
158                value: vec![
159                    EnumValueDescriptorProto {
160                        name: Some("MALE".to_string()),
161                        number: Some(0),
162                        ..Default::default()
163                    },
164                    EnumValueDescriptorProto {
165                        name: Some("FEMALE".to_string()),
166                        number: Some(1),
167                        ..Default::default()
168                    },
169                ],
170                ..Default::default()
171            }],
172            message_type: vec![DescriptorProto {
173                name: Some("Person".to_string()),
174                field: vec![
175                    FieldDescriptorProto {
176                        r#type: Some(Type::Int32.into()),
177                        name: Some("id".to_string()),
178                        number: Some(1),
179                        label: Some(Label::Optional.into()),
180                        ..Default::default()
181                    },
182                    FieldDescriptorProto {
183                        r#type: Some(Type::String.into()),
184                        name: Some("otherAttribute".to_string()),
185                        number: Some(2),
186                        label: Some(Label::Required.into()),
187                        ..Default::default()
188                    },
189                ],
190                ..Default::default()
191            }],
192            ..Default::default()
193        }
194    }
195
196    pub fn get_fds_encoded(files: Vec<FileDescriptorProto>) -> Vec<u8> {
197        let mut fds_encoded = Vec::new();
198
199        FileDescriptorSet { file: files }
200            .encode(&mut fds_encoded)
201            .unwrap();
202
203        fds_encoded
204    }
205}