prpc_build/

prost.rs

use super::{client, server, Attributes};
use proc_macro2::TokenStream;
use prost_build::{Config, Method, Service};
use quote::ToTokens;
use std::ffi::OsString;
use std::io;
use std::path::{Path, PathBuf};

/// Configure `prpc-build` code generation.
///
/// Use [`compile_protos`] instead if you don't need to tweak anything.
pub fn configure() -> Builder {
    Builder {
        build_client: true,
        build_server: true,
        build_scale_ext: true,
        out_dir: None,
        extern_path: Vec::new(),
        field_attributes: Vec::new(),
        type_attributes: Vec::new(),
        server_attributes: Attributes::default(),
        client_attributes: Attributes::default(),
        proto_path: "super".to_string(),
        compile_well_known_types: false,
        format: true,
        emit_package: true,
        emit_service_name: true,
        keep_service_names: Vec::new(),
        protoc_args: Vec::new(),
        file_descriptor_set_path: None,
        mod_prefix: Default::default(),
        type_prefix: Default::default(),
    }
}

/// Simple `.proto` compiling. Use [`configure`] instead if you need more options.
///
/// The include directory will be the parent folder of the specified path.
/// The package name will be the filename without the extension.
pub fn compile_protos(proto: impl AsRef<Path>) -> io::Result<()> {
    let proto_path: &Path = proto.as_ref();

    // directory the main .proto file resides in
    let proto_dir = proto_path
        .parent()
        .expect("proto file should reside in a directory");

    self::configure().compile(&[proto_path], &[proto_dir])?;

    Ok(())
}

impl crate::Service for Service {
    type Method = Method;
    type Comment = String;

    fn name(&self) -> &str {
        &self.name
    }

    fn package(&self) -> &str {
        &self.package
    }

    fn identifier(&self) -> &str {
        &self.proto_name
    }

    fn comment(&self) -> &[Self::Comment] {
        &self.comments.leading[..]
    }

    fn methods(&self) -> &[Self::Method] {
        &self.methods[..]
    }
}

impl crate::Method for Method {
    type Comment = String;

    fn name(&self) -> &str {
        &self.name
    }

    fn identifier(&self) -> &str {
        &self.proto_name
    }

    fn client_streaming(&self) -> bool {
        self.client_streaming
    }

    fn server_streaming(&self) -> bool {
        self.server_streaming
    }

    fn comment(&self) -> &[Self::Comment] {
        &self.comments.leading[..]
    }

    fn request_response_name(
        &self,
        proto_path: &str,
        compile_well_known_types: bool,
    ) -> (Option<TokenStream>, TokenStream) {
        let request = if is_empty_type(&self.input_proto_type) {
            None
        } else {
            Some(
                if (is_google_type(&self.input_proto_type) && !compile_well_known_types)
                    || self.input_type.starts_with("::")
                {
                    self.input_type.parse::<TokenStream>().unwrap()
                } else if self.input_type.starts_with("crate::") {
                    syn::parse_str::<syn::Path>(&self.input_type)
                        .unwrap()
                        .to_token_stream()
                } else {
                    syn::parse_str::<syn::Path>(&format!("{proto_path}::{}", self.input_type))
                        .unwrap()
                        .to_token_stream()
                },
            )
        };

        let response = if (is_google_type(&self.output_proto_type) && !compile_well_known_types)
            || self.output_type.starts_with("::")
        {
            self.output_type.parse::<TokenStream>().unwrap()
        } else if self.output_type.starts_with("crate::") {
            syn::parse_str::<syn::Path>(&self.output_type)
                .unwrap()
                .to_token_stream()
        } else {
            syn::parse_str::<syn::Path>(&format!("{proto_path}::{}", self.output_type))
                .unwrap()
                .to_token_stream()
        };

        (request, response)
    }
}

fn is_google_type(ty: &str) -> bool {
    ty.starts_with(".google.protobuf")
}

fn is_empty_type(ty: &str) -> bool {
    ty == ".google.protobuf.Empty"
}

struct ServiceGenerator {
    builder: Builder,
    clients: TokenStream,
    servers: TokenStream,
}

impl ServiceGenerator {
    fn new(builder: Builder) -> Self {
        ServiceGenerator {
            builder,
            clients: TokenStream::default(),
            servers: TokenStream::default(),
        }
    }
}

impl prost_build::ServiceGenerator for ServiceGenerator {
    fn generate(&mut self, service: prost_build::Service, _buf: &mut String) {
        if self.builder.build_server {
            let server = server::generate(&service, &self.builder);
            self.servers.extend(server);
        }

        if self.builder.build_client {
            let client = client::generate(&service, &self.builder);
            self.clients.extend(client);
        }
    }

    fn finalize(&mut self, buf: &mut String) {
        if self.builder.build_client && !self.clients.is_empty() {
            let code = format!("{}", self.clients);
            buf.push_str(&code);

            self.clients = TokenStream::default();
        }

        if self.builder.build_server && !self.servers.is_empty() {
            let code = format!("{}", self.servers);
            buf.push_str(&code);

            self.servers = TokenStream::default();
        }
    }
}

/// Service generator builder.
#[derive(Debug, Clone)]
pub struct Builder {
    pub(crate) build_client: bool,
    pub(crate) build_server: bool,
    pub(crate) build_scale_ext: bool,
    pub(crate) extern_path: Vec<(String, String)>,
    pub(crate) field_attributes: Vec<(String, String)>,
    pub(crate) type_attributes: Vec<(String, String)>,
    pub(crate) server_attributes: Attributes,
    pub(crate) client_attributes: Attributes,
    pub(crate) proto_path: String,
    pub(crate) emit_package: bool,
    pub(crate) emit_service_name: bool,
    pub(crate) keep_service_names: Vec<String>,
    pub(crate) compile_well_known_types: bool,
    pub(crate) protoc_args: Vec<OsString>,

    mod_prefix: String,
    type_prefix: String,
    file_descriptor_set_path: Option<PathBuf>,
    out_dir: Option<PathBuf>,
    format: bool,
}

impl Builder {
    /// Enable or disable client code generation.
    pub fn build_client(mut self, enable: bool) -> Self {
        self.build_client = enable;
        self
    }

    /// Enable or disable server code generation.
    pub fn build_server(mut self, enable: bool) -> Self {
        self.build_server = enable;
        self
    }

    /// Enable or disable scale codec extensions generation.
    pub fn build_scale_ext(mut self, enable: bool) -> Self {
        self.build_scale_ext = enable;
        self
    }

    /// Enable the output to be formated by rustfmt.
    pub fn format(mut self, run: bool) -> Self {
        self.format = run;
        self
    }

    /// Module prefix of the generated code.
    pub fn mod_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.mod_prefix = prefix.into();
        self
    }

    /// Type prefix of the scale codec anotation in the proto file.
    pub fn type_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.type_prefix = prefix.into();
        self
    }

    /// Set the output directory to generate code to.
    ///
    /// Defaults to the `OUT_DIR` environment variable.
    pub fn out_dir(mut self, out_dir: impl AsRef<Path>) -> Self {
        self.out_dir = Some(out_dir.as_ref().to_path_buf());
        self
    }

    /// Declare externally provided Protobuf package or type.
    ///
    /// Passed directly to `prost_build::Config.extern_path`.
    /// Note that both the Protobuf path and the rust package paths should both be fully qualified.
    /// i.e. Protobuf paths should start with "." and rust paths should start with "::"
    pub fn extern_path(mut self, proto_path: impl AsRef<str>, rust_path: impl AsRef<str>) -> Self {
        self.extern_path.push((
            proto_path.as_ref().to_string(),
            rust_path.as_ref().to_string(),
        ));
        self
    }

    /// Add additional attribute to matched messages, enums, and one-offs.
    ///
    /// Passed directly to `prost_build::Config.field_attribute`.
    pub fn field_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
        self.field_attributes
            .push((path.as_ref().to_string(), attribute.as_ref().to_string()));
        self
    }

    /// Add additional attribute to matched messages, enums, and one-offs.
    ///
    /// Passed directly to `prost_build::Config.type_attribute`.
    pub fn type_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
        self.type_attributes
            .push((path.as_ref().to_string(), attribute.as_ref().to_string()));
        self
    }

    /// Add additional attribute to matched server `mod`s. Matches on the package name.
    pub fn server_mod_attribute<P: AsRef<str>, A: AsRef<str>>(
        mut self,
        path: P,
        attribute: A,
    ) -> Self {
        self.server_attributes
            .push_mod(path.as_ref().to_string(), attribute.as_ref().to_string());
        self
    }

    /// Add additional attribute to matched service servers. Matches on the service name.
    pub fn server_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
        self.server_attributes
            .push_struct(path.as_ref().to_string(), attribute.as_ref().to_string());
        self
    }

    /// Add additional attribute to matched client `mod`s. Matches on the package name.
    pub fn client_mod_attribute<P: AsRef<str>, A: AsRef<str>>(
        mut self,
        path: P,
        attribute: A,
    ) -> Self {
        self.client_attributes
            .push_mod(path.as_ref().to_string(), attribute.as_ref().to_string());
        self
    }

    /// Add additional attribute to matched service clients. Matches on the service name.
    pub fn client_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
        self.client_attributes
            .push_struct(path.as_ref().to_string(), attribute.as_ref().to_string());
        self
    }

    /// Set the path to where to search for the Request/Response proto structs
    /// live relative to the module where you call `include_proto!`.
    ///
    /// This defaults to `super` since we will generate code in a module.
    pub fn proto_path(mut self, proto_path: impl AsRef<str>) -> Self {
        self.proto_path = proto_path.as_ref().to_string();
        self
    }

    /// Configure Prost `protoc_args` build arguments.
    ///
    /// Note: Enabling `--experimental_allow_proto3_optional` requires protobuf >= 3.12.
    pub fn protoc_arg<A: AsRef<str>>(mut self, arg: A) -> Self {
        self.protoc_args.push(arg.as_ref().into());
        self
    }

    /// Emits RPC endpoints with no attached package. Effectively ignores protofile package declaration from rpc context.
    ///
    /// This effectively sets prost's exported package to an empty string.
    pub fn disable_package_emission(mut self) -> Self {
        self.emit_package = false;
        self
    }

    /// Disable emitting service name in rpc endpoints.
    pub fn disable_service_name_emission(mut self) -> Self {
        self.emit_service_name = false;
        self
    }

    /// Keep the service name in rpc endpoints.
    pub fn keep_service_name(mut self, name: impl Into<String>) -> Self {
        self.keep_service_names.push(name.into());
        self
    }

    /// Enable or disable directing Prost to compile well-known protobuf types instead
    /// of using the already-compiled versions available in the `prost-types` crate.
    ///
    /// This defaults to `false`.
    pub fn compile_well_known_types(mut self, compile_well_known_types: bool) -> Self {
        self.compile_well_known_types = compile_well_known_types;
        self
    }

    /// When set, the `FileDescriptorSet` generated by `protoc` is written to the provided
    /// filesystem path.
    ///
    /// This option can be used in conjunction with the [`include_bytes!`] macro and the types in
    /// the `prost-types` crate for implementing reflection capabilities, among other things.
    pub fn file_descriptor_set_path(mut self, path: impl Into<PathBuf>) -> Self {
        self.file_descriptor_set_path = Some(path.into());
        self
    }

    /// Compile the .proto files and execute code generation.
    pub fn compile(
        self,
        protos: &[impl AsRef<Path>],
        includes: &[impl AsRef<Path>],
    ) -> io::Result<()> {
        self.compile_with_config(Config::new(), protos, includes)
    }

    /// Compile the .proto files and execute code generation using a
    /// custom `prost_build::Config`.
    pub fn compile_with_config(
        self,
        mut config: Config,
        protos: &[impl AsRef<Path>],
        includes: &[impl AsRef<Path>],
    ) -> io::Result<()> {
        let out_dir = if let Some(out_dir) = self.out_dir.as_ref() {
            out_dir.clone()
        } else {
            PathBuf::from(std::env::var("OUT_DIR").unwrap())
        };
        if !out_dir.exists() {
            fs_err::create_dir_all(&out_dir)?;
        }

        let format = self.format;

        config.out_dir(out_dir.clone());
        for (proto_path, rust_path) in self.extern_path.iter() {
            config.extern_path(proto_path, rust_path);
        }
        for (prost_path, attr) in self.field_attributes.iter() {
            config.field_attribute(prost_path, attr);
        }
        for (prost_path, attr) in self.type_attributes.iter() {
            config.type_attribute(prost_path, attr);
        }
        if self.compile_well_known_types {
            config.compile_well_known_types();
        }

        for arg in self.protoc_args.iter() {
            config.protoc_arg(arg);
        }

        let file_descriptor_set_path =
            if let Some(file_descriptor_set_path) = &self.file_descriptor_set_path {
                file_descriptor_set_path.clone()
            } else {
                out_dir.join("file_descriptor_set.bin")
            };
        config.file_descriptor_set_path(file_descriptor_set_path.clone());

        config.service_generator(Box::new(ServiceGenerator::new(self.clone())));

        let protoc = match std::env::var("PROTOC") {
            Ok(path) => PathBuf::from(path),
            Err(_) => protoc::protoc(),
        };
        config.protoc_executable(protoc);
        config.compile_protos(protos, includes)?;

        if self.build_scale_ext {
            let patch_file = out_dir.join("protos_codec_extensions.rs");
            crate::protos_codec_extension::extend_types(
                &file_descriptor_set_path,
                patch_file,
                &self.mod_prefix,
                &self.type_prefix,
            );
        }

        {
            if format {
                super::fmt(out_dir.to_str().expect("expected utf8 out_dir"));
            }
        }

        Ok(())
    }

    /// Enable serde serialization/deserialization for all generated types.
    /// This adds the necessary derives and attributes for serde compatibility.
    pub fn enable_serde_extension(self) -> Self {
        self.type_attribute(".", "#[::prpc::serde_helpers::prpc_serde_bytes]")
            .type_attribute(".", "#[derive(::serde::Serialize, ::serde::Deserialize)]")
            .field_attribute(".", "#[serde(default)]")
    }

    /// Compile all .proto files in the specified directory.
    /// The include directory will be the same as the proto directory.
    pub fn compile_dir(self, proto_dir: impl AsRef<Path>) -> io::Result<()> {
        let proto_dir = proto_dir.as_ref();

        let proto_files: Vec<PathBuf> = std::fs::read_dir(proto_dir)?
            .filter_map(|entry| {
                let entry = entry.ok()?;
                let path = entry.path();
                if path.extension()?.to_str()? == "proto" {
                    Some(path)
                } else {
                    None
                }
            })
            .collect();

        if proto_files.is_empty() {
            return Ok(());
        }

        self.compile(&proto_files, &[proto_dir])
    }
}