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 pub fn new() -> Self {
29 Self::default()
30 }
31
32 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 pub fn generate_enum_values(&mut self) -> &mut Self {
46 self.generate_enum_values = true;
47 self
48 }
49
50 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 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 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}