prost_wkt_build/
lib.rs

1use heck::ToUpperCamelCase;
2use quote::{format_ident, quote};
3use std::fs::{File, OpenOptions};
4use std::io::Write;
5use std::path::PathBuf;
6
7pub use prost::Message;
8pub use prost_types::FileDescriptorSet;
9
10use prost_build::Module;
11
12pub fn add_serde(out: PathBuf, descriptor: FileDescriptorSet) {
13    for fd in &descriptor.file {
14        let package_name = match fd.package {
15            Some(ref pkg) => pkg,
16            None => continue,
17        };
18
19        let rust_path = out
20            .join(Module::from_protobuf_package_name(package_name).to_file_name_or(package_name));
21
22        // In some cases the generated file would be in empty. These files are no longer created by Prost, so
23        // we'll create here. Otherwise we append.
24        let mut rust_file = OpenOptions::new()
25            .create(true)
26            .append(true)
27            .open(rust_path)
28            .unwrap();
29
30        for msg in &fd.message_type {
31            let message_name = match msg.name {
32                Some(ref name) => name,
33                None => continue,
34            };
35
36            let type_url = format!("type.googleapis.com/{package_name}.{message_name}");
37
38            gen_trait_impl(&mut rust_file, package_name, message_name, &type_url);
39        }
40    }
41}
42
43// This method uses the `heck` crate (the same that prost uses) to properly format the message name
44// to UpperCamelCase as the prost_build::ident::{to_snake, to_upper_camel} methods
45// in the `ident` module of prost_build is private.
46fn gen_trait_impl(rust_file: &mut File, package_name: &str, message_name: &str, type_url: &str) {
47    let type_name = message_name.to_upper_camel_case();
48    let type_name = format_ident!("{}", type_name);
49
50    let tokens = quote! {
51        #[allow(dead_code)]
52        const _: () = {
53            use ::prost_wkt::typetag;
54            #[typetag::serde(name=#type_url)]
55            impl ::prost_wkt::MessageSerde for #type_name {
56                fn package_name(&self) -> &'static str {
57                    #package_name
58                }
59                fn message_name(&self) -> &'static str {
60                    #message_name
61                }
62                fn type_url(&self) -> &'static str {
63                    #type_url
64                }
65                fn new_instance(&self, data: Vec<u8>) -> ::std::result::Result<Box<dyn ::prost_wkt::MessageSerde>, ::prost::DecodeError> {
66                    let mut target = Self::default();
67                    ::prost::Message::merge(&mut target, data.as_slice())?;
68                    let erased: ::std::boxed::Box<dyn ::prost_wkt::MessageSerde> = ::std::boxed::Box::new(target);
69                    Ok(erased)
70                }
71                fn try_encoded(&self) -> ::std::result::Result<::std::vec::Vec<u8>, ::prost::EncodeError> {
72                    let mut buf = ::std::vec::Vec::with_capacity(::prost::Message::encoded_len(self));
73                    ::prost::Message::encode(self, &mut buf)?;
74                    Ok(buf)
75                }
76            }
77
78            ::prost_wkt::inventory::submit!{
79                ::prost_wkt::MessageSerdeDecoderEntry {
80                    type_url: #type_url,
81                    decoder: |buf: &[u8]| {
82                        let msg: #type_name = ::prost::Message::decode(buf)?;
83                        Ok(::std::boxed::Box::new(msg))
84                    }
85                }
86            }
87
88            impl ::prost::Name for #type_name {
89                const PACKAGE: &'static str = #package_name;
90                const NAME: &'static str = #message_name;
91
92                fn type_url() -> String {
93                    #type_url.to_string()
94                }
95            }
96        };
97    };
98
99    writeln!(rust_file).unwrap();
100    writeln!(rust_file, "{}", &tokens).unwrap();
101}
prost_wkt_build/lib.rs

prost_wkt_build/
lib.rs
