use crate::ir::StructIR;
use proc_macro2::TokenStream;
use quote::quote;
pub fn generate_trait_impl(ir: &StructIR, config: &super::MacroConfig) -> TokenStream {
let struct_name = &ir.name;
let to_df_trait = &config.traits.to_dataframe;
let columnar_trait = &config.traits.columnar;
let pp = config.external_paths.prelude();
let (impl_generics, ty_generics, where_clause) =
super::bounds::impl_parts_with_bounds(ir, config);
if ir.columns.is_empty() {
return quote! {
#[automatically_derived]
impl #impl_generics #to_df_trait for #struct_name #ty_generics #where_clause {
fn to_dataframe(&self) -> #pp::PolarsResult<#pp::DataFrame> {
<Self as #columnar_trait>::columnar_from_refs(&[self])
}
fn empty_dataframe() -> #pp::PolarsResult<#pp::DataFrame> {
#pp::DataFrame::new_infer_height(::std::vec![])
}
fn schema() -> #pp::PolarsResult<::std::vec::Vec<(::std::string::String, #pp::DataType)>> {
::std::result::Result::Ok(::std::vec::Vec::new())
}
}
};
}
let empty_series_creations: Vec<TokenStream> = ir
.columns
.iter()
.map(|column| super::schema::build_empty_series(column, config))
.collect();
let schema_entries: Vec<TokenStream> = ir
.columns
.iter()
.map(|column| super::schema::build_schema_entries(column, config))
.collect();
quote! {
#[automatically_derived]
impl #impl_generics #to_df_trait for #struct_name #ty_generics #where_clause {
fn to_dataframe(&self) -> #pp::PolarsResult<#pp::DataFrame> {
<Self as #columnar_trait>::columnar_from_refs(&[self])
}
fn empty_dataframe() -> #pp::PolarsResult<#pp::DataFrame> {
let mut all_series: ::std::vec::Vec<#pp::Column> = ::std::vec::Vec::new();
#(
all_series.extend(#empty_series_creations);
)*
#pp::DataFrame::new_infer_height(all_series)
}
fn schema() -> #pp::PolarsResult<::std::vec::Vec<(::std::string::String, #pp::DataType)>> {
let mut fields: ::std::vec::Vec<(::std::string::String, #pp::DataType)> = ::std::vec::Vec::new();
#(
fields.extend(#schema_entries);
)*
::std::result::Result::Ok(fields)
}
}
}
}