pyo3_tonic/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::str;
4
5use prost::Message;
6use prost_types::compiler::CodeGeneratorRequest;
7use protoc_gen_prost::{Generator, InvalidParameter, ModuleRequestSet, Param, Params};
8use tonic_build::Attributes;
9
10use self::{generator::TonicGenerator, resolver::Resolver};
11
12mod generator;
13mod resolver;
14mod util;
15
16/// Execute the core _Prost!_ generator from a raw [`CodeGeneratorRequest`]
17pub fn execute(raw_request: &[u8]) -> protoc_gen_prost::Result {
18    let request = CodeGeneratorRequest::decode(raw_request)?;
19    let params = request.parameter().parse::<Parameters>()?;
20
21    let module_request_set = ModuleRequestSet::new(
22        request.file_to_generate,
23        request.proto_file,
24        raw_request,
25        params.default_package_filename.as_deref(),
26    )?;
27
28    let resolver = Resolver::new(params.extern_path, params.compile_well_known_types);
29    let mut generator = TonicGenerator {
30        resolver,
31        generate_server: !params.no_server,
32        generate_client: !params.no_client,
33        generate_transport: !params.no_transport,
34        server_attributes: params.server_attributes,
35        client_attributes: params.client_attributes,
36        emit_package: !params.disable_package_emission,
37        insert_include: !params.no_include,
38    };
39
40    let files = generator.generate(&module_request_set)?;
41
42    Ok(files)
43}
44
45/// Parameters use to configure [`Generator`]s built into `protoc-gen-prost-serde`
46///
47/// [`Generator`]: protoc_gen_prost::generators::Generator
48#[derive(Debug, Default)]
49struct Parameters {
50    default_package_filename: Option<String>,
51    extern_path: Vec<(String, String)>,
52    server_attributes: Attributes,
53    client_attributes: Attributes,
54    compile_well_known_types: bool,
55    disable_package_emission: bool,
56    no_server: bool,
57    no_client: bool,
58    no_transport: bool,
59    no_include: bool,
60}
61
62impl str::FromStr for Parameters {
63    type Err = InvalidParameter;
64    fn from_str(s: &str) -> Result<Self, Self::Err> {
65        let mut ret_val = Self::default();
66        for param in Params::from_protoc_plugin_opts(s)? {
67            match param {
68                Param::Parameter {
69                    param: "default_package_filename",
70                }
71                | Param::Value {
72                    param: "default_package_filename",
73                    ..
74                } => ret_val.default_package_filename = param.value().map(|s| s.to_string()),
75                Param::KeyValue {
76                    param: "extern_path",
77                    key: prefix,
78                    value: module,
79                } => ret_val.extern_path.push((prefix.to_string(), module)),
80                Param::Parameter {
81                    param: "compile_well_known_types",
82                }
83                | Param::Value {
84                    param: "compile_well_known_types",
85                    value: "true",
86                } => ret_val.compile_well_known_types = true,
87                Param::Value {
88                    param: "compile_well_known_types",
89                    value: "false",
90                } => (),
91                Param::Parameter {
92                    param: "disable_package_emission",
93                }
94                | Param::Value {
95                    param: "disable_package_emission",
96                    value: "true",
97                } => ret_val.disable_package_emission = true,
98                Param::Value {
99                    param: "disable_package_emission",
100                    value: "false",
101                } => (),
102                Param::Parameter { param: "no_server" }
103                | Param::Value {
104                    param: "no_server",
105                    value: "true",
106                } => ret_val.no_server = true,
107                Param::Value {
108                    param: "no_server",
109                    value: "false",
110                } => (),
111                Param::Parameter { param: "no_client" }
112                | Param::Value {
113                    param: "no_client",
114                    value: "true",
115                } => ret_val.no_client = true,
116                Param::Value {
117                    param: "no_client",
118                    value: "false",
119                } => (),
120                Param::Parameter {
121                    param: "no_transport",
122                }
123                | Param::Value {
124                    param: "no_transport",
125                    value: "true",
126                } => ret_val.no_transport = true,
127                Param::Value {
128                    param: "no_transport",
129                    value: "false",
130                } => (),
131                Param::Parameter {
132                    param: "no_include",
133                }
134                | Param::Value {
135                    param: "no_include",
136                    value: "true",
137                } => ret_val.no_include = true,
138                Param::Value {
139                    param: "no_include",
140                    value: "false",
141                } => (),
142                Param::KeyValue {
143                    param: "client_mod_attribute",
144                    key: prefix,
145                    value: attribute,
146                } => ret_val
147                    .client_attributes
148                    .push_mod(prefix, attribute.replace(r"\,", ",").replace(r"\\", r"\")),
149                Param::KeyValue {
150                    param: "client_attribute",
151                    key: prefix,
152                    value: attribute,
153                } => ret_val
154                    .client_attributes
155                    .push_struct(prefix, attribute.replace(r"\,", ",").replace(r"\\", r"\")),
156                Param::KeyValue {
157                    param: "server_mod_attribute",
158                    key: prefix,
159                    value: attribute,
160                } => ret_val
161                    .server_attributes
162                    .push_mod(prefix, attribute.replace(r"\,", ",").replace(r"\\", r"\")),
163                Param::KeyValue {
164                    param: "server_attribute",
165                    key: prefix,
166                    value: attribute,
167                } => ret_val
168                    .server_attributes
169                    .push_struct(prefix, attribute.replace(r"\,", ",").replace(r"\\", r"\")),
170                _ => return Err(InvalidParameter::from(param)),
171            }
172        }
173
174        Ok(ret_val)
175    }
176}