olai-codegen 0.0.3

//! PyO3 binding generation for protobuf-defined services.
//!
//! The dispatch and "which methods land where" logic is shared with the Node emitter via
//! [`crate::codegen::bindings`]; this module supplies the PyO3-specific scaffold, marshaling, and
//! method bodies through [`PyBackend`]'s [`BindingBackend`] impl.

use itertools::Itertools;
use proc_macro2::{Ident, TokenStream};
use quote::{format_ident, quote};

use super::super::format_tokens;
use crate::codegen::bindings::builder::{OptionalSetter, SetterRender, generate_builder_pattern};
use crate::codegen::bindings::{BindingBackend, ShapeParts, driver};
use crate::codegen::{BindingMode, MethodHandler, ServiceHandler};
use crate::parsing::types::{BaseType, RenderContext};

/// PyO3 binding backend. Holds the configured Python error/result type idents so the method bodies
/// can wrap calls in the right `PyResult`/error conversion.
pub(crate) struct PyBackend {
    py_error_type: Ident,
    py_result_type: Ident,
}

impl PyBackend {
    fn from_service(service: &ServiceHandler<'_>) -> Self {
        let bindings = service
            .config
            .bindings
            .as_ref()
            .expect("bindings config required for python output");
        Self {
            py_error_type: format_ident!("{}", bindings.py_error_type),
            py_result_type: format_ident!("{}", bindings.py_result_type),
        }
    }

    fn from_services(services: &[ServiceHandler<'_>]) -> Self {
        let bindings = services
            .first()
            .and_then(|s| s.config.bindings.as_ref())
            .expect("bindings config required for python output");
        Self {
            py_error_type: format_ident!("{}", bindings.py_error_type),
            py_result_type: format_ident!("{}", bindings.py_result_type),
        }
    }
}

pub fn main_module(services: &[ServiceHandler<'_>]) -> crate::error::Result<String> {
    PyBackend::from_services(services).main_module(services)
}

pub(crate) fn generate(service: &ServiceHandler<'_>) -> crate::error::Result<String> {
    PyBackend::from_service(service).generate_service(service)
}

impl BindingBackend for PyBackend {
    fn generate_service(&self, service: &ServiceHandler<'_>) -> crate::error::Result<String> {
        let bindings = service
            .config
            .bindings
            .as_ref()
            .expect("bindings config required for python output");

        let rust_client_ident = service.client_type();
        let client_ident = format_ident!("Py{}", rust_client_ident);
        let rust_client_name = rust_client_ident.to_string();

        let client_crate = format_ident!("{}", bindings.client_crate_name);
        let py_error_type = &self.py_error_type;
        let py_result_type = &self.py_result_type;

        let methods = driver::scoped_client_methods(self, service);
        let mod_path = service.models_path();

        let tokens = quote! {
            use std::collections::HashMap;
            use pyo3::prelude::*;
            use #client_crate::#rust_client_ident;
            use #mod_path::*;
            use crate::error::{#py_error_type, #py_result_type};
            use crate::runtime::get_runtime;

            #[pyclass(name = #rust_client_name)]
            pub struct #client_ident {
                pub(crate) client: #rust_client_ident,
            }

            #[pymethods]
            impl #client_ident {
                #(#methods)*
            }

            impl #client_ident {
                pub fn new(client: #rust_client_ident) -> Self {
                    Self { client }
                }
            }
        };

        format_tokens(tokens)
    }

    fn main_module(&self, services: &[ServiceHandler<'_>]) -> crate::error::Result<String> {
        // Only resource-scoped services have a per-service scoped module; resource-less services'
        // methods live on the root client.
        let service_modules = services.iter().filter(|s| s.is_resource_scoped()).map(|s| {
            let module_name = format_ident!("{}", s.plan.base_path);
            quote! { pub mod #module_name; }
        });

        let bindings = services
            .first()
            .and_then(|s| s.config.bindings.as_ref())
            .expect("bindings config required for python output");

        let aggregate_client_name = &bindings.aggregate_client_name;
        let client_crate = format_ident!("{}", bindings.client_crate_name);
        let aggregate_client_ident = format_ident!("{}", aggregate_client_name);
        let py_aggregate_client_ident = format_ident!("Py{}", aggregate_client_name);
        let py_error_type = &self.py_error_type;
        let py_result_type = &self.py_result_type;

        let mut sorted_services = services.iter().collect_vec();
        sorted_services.sort_by(|a, b| a.plan.service_name.cmp(&b.plan.service_name));

        let mod_paths = sorted_services.iter().map(|s| {
            let mod_path = s.models_path();
            quote! { use #mod_path::*; }
        });

        // Only resource-scoped services expose a per-service scoped client to import.
        let codegen_imports = sorted_services
            .iter()
            .filter(|s| s.is_resource_scoped())
            .map(|s| {
                let mod_name = format_ident!("{}", s.plan.base_path);
                let client_name = format_ident!("Py{}", s.client_type().to_string());
                quote! { use crate::codegen::#mod_name::#client_name; }
            });

        let methods = driver::root_client_methods(self, &sorted_services);
        let resource_accessor_methods = driver::resource_accessor_methods(self, &sorted_services);

        let tokens = quote! {
            #(#service_modules)*

            use std::collections::HashMap;
            use futures::stream::TryStreamExt;
            use pyo3::prelude::*;
            use #client_crate::{#aggregate_client_ident};
            use crate::error::{#py_error_type, #py_result_type};
            use crate::runtime::get_runtime;
            #(#mod_paths)*
            #(#codegen_imports)*

            #[pyclass(name = #aggregate_client_name)]
            pub struct #py_aggregate_client_ident {
                client: #aggregate_client_ident
            }

            #[pymethods]
            impl #py_aggregate_client_ident {
                #[new]
                #[pyo3(signature = (base_url, token = None))]
                pub fn new(base_url: String, token: Option<String>) -> PyResult<Self> {
                    let client = if let Some(token) = token {
                        olai_http::CloudClient::new_with_token(token)
                    } else {
                        olai_http::CloudClient::new_unauthenticated()
                    };
                    let base_url = base_url.parse().map_err(#py_error_type::from)?;
                    Ok(Self { client: #aggregate_client_ident::new(client, base_url) })
                }

                #(#methods)*

                #(#resource_accessor_methods)*
            }
        };

        format_tokens(tokens)
    }

    fn param_defs(&self, method: &MethodHandler<'_>, mode: BindingMode) -> Vec<TokenStream> {
        python_param_defs(method, mode)
    }

    fn client_call(&self, method: &MethodHandler<'_>, mode: BindingMode) -> TokenStream {
        let (method_name, drop_path) = match mode {
            BindingMode::Scoped => (method.plan.resource_client_method(), true),
            BindingMode::Flat => (method.plan.base_method_ident(), false),
        };
        let args = method
            .required_parameters()
            .filter(|param| !(drop_path && param.is_path_param()))
            .map(|param| {
                let param_name = param.field_ident();
                quote! { #param_name }
            });
        quote! {
            self.client.#method_name(#(#args,)*)
        }
    }

    fn emit_list(&self, parts: &ShapeParts<'_>) -> TokenStream {
        let method = parts.method;
        let py_error_type = &self.py_error_type;
        let py_result_type = &self.py_result_type;
        let method_name = method.binding_method_name();
        let param_defs = &parts.param_defs;
        let signature = pyo3_signature(method, parts.mode);
        let client_call = &parts.client_call;
        let builder_calls = generate_builder_pattern(method, true, self);

        let items_field = method
            .list_output_field()
            .expect("List-shaped method must have a list output field");
        let response_type = method.field_type(&items_field.unified_type, RenderContext::ReturnType);

        quote! {
            #signature
            pub fn #method_name(
                &self,
                py: Python,
                #(#param_defs,)*
            ) -> #py_result_type<#response_type> {
                let mut request = #client_call;
                #(#builder_calls)*
                let runtime = get_runtime(py)?;
                py.allow_threads(|| {
                    let result = runtime.block_on(async move { request.into_stream().try_collect().await })?;
                    Ok::<_, #py_error_type>(result)
                })
            }
        }
    }

    fn emit_create(&self, parts: &ShapeParts<'_>) -> TokenStream {
        let method = parts.method;
        let py_error_type = &self.py_error_type;
        let py_result_type = &self.py_result_type;
        let method_name = method.binding_method_name();
        let response_type = method.output_type_or_unit();
        let param_defs = &parts.param_defs;
        let signature = pyo3_signature(method, parts.mode);
        let client_call = &parts.client_call;
        let builder_calls = generate_builder_pattern(method, false, self);

        quote! {
            #signature
            pub fn #method_name(
                &self,
                py: Python,
                #(#param_defs,)*
            ) -> #py_result_type<#response_type> {
                let mut request = #client_call;
                #(#builder_calls)*
                let runtime = get_runtime(py)?;
                py.allow_threads(|| {
                    let result = runtime.block_on(request.into_future())?;
                    Ok::<_, #py_error_type>(result)
                })
            }
        }
    }

    fn emit_get_update(&self, parts: &ShapeParts<'_>) -> TokenStream {
        let method = parts.method;
        let py_error_type = &self.py_error_type;
        let py_result_type = &self.py_result_type;
        let method_name = match parts.mode {
            BindingMode::Scoped => method.plan.resource_client_method(),
            BindingMode::Flat => method.binding_method_name(),
        };
        let response_type = method.output_type_or_unit();
        let param_defs = &parts.param_defs;
        let signature = pyo3_signature(method, parts.mode);
        let client_call = &parts.client_call;
        let builder_calls = generate_builder_pattern(method, false, self);

        quote! {
            #signature
            pub fn #method_name(
                &self,
                py: Python,
                #(#param_defs,)*
            ) -> #py_result_type<#response_type> {
                let mut request = #client_call;
                #(#builder_calls)*
                let runtime = get_runtime(py)?;
                py.allow_threads(|| {
                    let result = runtime.block_on(request.into_future())?;
                    Ok::<_, #py_error_type>(result)
                })
            }
        }
    }

    fn emit_delete(&self, parts: &ShapeParts<'_>) -> TokenStream {
        let method = parts.method;
        let py_error_type = &self.py_error_type;
        let py_result_type = &self.py_result_type;
        let method_name = match parts.mode {
            BindingMode::Scoped => method.plan.resource_client_method(),
            BindingMode::Flat => method.binding_method_name(),
        };
        let param_defs = &parts.param_defs;
        let signature = pyo3_signature(method, parts.mode);
        let client_call = &parts.client_call;
        let builder_calls = generate_builder_pattern(method, false, self);

        quote! {
            #signature
            pub fn #method_name(
                &self,
                py: Python,
                #(#param_defs,)*
            ) -> #py_result_type<()> {
                let mut request = #client_call;
                #(#builder_calls)*
                let runtime = get_runtime(py)?;
                py.allow_threads(|| {
                    runtime.block_on(request.into_future())?;
                    Ok::<_, #py_error_type>(())
                })
            }
        }
    }

    fn emit_resource_accessor(&self, service: &ServiceHandler<'_>) -> Option<TokenStream> {
        let spec = service.accessor_spec()?;
        let method_name = format_ident!("{}", spec.singular);
        let client_name = format_ident!("Py{}", service.client_type().to_string());

        let param_idents: Vec<_> = spec.params.iter().map(|p| format_ident!("{}", p)).collect();
        let param_list = param_idents
            .iter()
            .map(|id| quote! { #id: String })
            .collect::<Vec<_>>();

        // A nested resource has a composite `full_name`; the underlying Rust client accessor takes a
        // single `full_name` string, so join the decomposed Python params back together with `.`.
        let method_call = if spec.nested {
            let from_full_name_method = format_ident!("{}_from_full_name", spec.singular);
            let format_str = spec.join_format();
            quote! {
                pub fn #method_name(&self, #(#param_list),*) -> #client_name {
                    let full_name = format!(#format_str, #(#param_idents),*);
                    #client_name {
                        client: self.client.#from_full_name_method(full_name),
                    }
                }
            }
        } else {
            let param_refs = param_idents
                .iter()
                .map(|id| quote! { #id })
                .collect::<Vec<_>>();
            quote! {
                pub fn #method_name(&self, #(#param_list),*) -> #client_name {
                    #client_name {
                        client: self.client.#method_name(#(#param_refs),*),
                    }
                }
            }
        };

        Some(method_call)
    }
}

impl SetterRender for PyBackend {
    fn supports(&self, _setter: &OptionalSetter<'_>) -> bool {
        // PyO3 accepts every parameter type; no capability filter.
        true
    }

    fn render_setter(
        &self,
        setter: &OptionalSetter<'_>,
        param_ident: &Ident,
        with_method: &Ident,
    ) -> TokenStream {
        // Maps and repeated fields are `Option`-wrapped at the binding boundary; guard before
        // forwarding. Everything else is forwarded directly.
        if matches!(setter.field_type.base_type, BaseType::Map(_, _))
            || setter.field_type.is_repeated
        {
            quote! {
                if let Some(#param_ident) = #param_ident {
                    request = request.#with_method(#param_ident);
                }
            }
        } else {
            quote! {
                request = request.#with_method(#param_ident);
            }
        }
    }
}

/// Build the Python method signature param defs from the shared [`MethodHandler::param_plan`], so
/// path/`page_token` filtering lives in one place.
fn python_param_defs(method: &MethodHandler<'_>, mode: BindingMode) -> Vec<TokenStream> {
    method
        .param_plan(mode)
        .into_iter()
        .map(|p| {
            let param_name = p.field_ident();
            let rust_type = method.field_type(p.field_type(), RenderContext::PythonParameter);
            quote! { #param_name: #rust_type }
        })
        .collect()
}

/// The `#[pyo3(signature = (...))]` attribute for a method, derived from its param plan.
fn pyo3_signature(method: &MethodHandler<'_>, mode: BindingMode) -> TokenStream {
    let signature_parts = method
        .param_plan(mode)
        .iter()
        .map(|p| {
            if p.is_optional() {
                format!("{} = None", p.name())
            } else {
                p.name().to_string()
            }
        })
        .collect_vec();
    if signature_parts.is_empty() {
        quote! {}
    } else {
        let signature_string = signature_parts.join(", ");
        // signature_parts are protobuf field names (always valid identifiers), so parsing
        // as a TokenStream is infallible.
        let tokens = signature_string
            .parse::<proc_macro2::TokenStream>()
            .unwrap();
        quote! {
            #[pyo3(signature = (#tokens))]
        }
    }
}