#![warn(missing_docs)]
#![warn(rustdoc::missing_crate_level_docs)]
use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::{
Data, DeriveInput, Fields, GenericArgument, PathArguments, Type, TypePath, parse_macro_input,
};
#[proc_macro_derive(ToDataFrame)]
pub fn derive_to_dataframe(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let name = &input.ident;
let fields = match &input.data {
Data::Struct(data) => match &data.fields {
Fields::Named(fields) => &fields.named,
_ => {
return syn::Error::new_spanned(
&input,
"ToDataFrame only supports structs with named fields",
)
.to_compile_error()
.into();
}
},
_ => {
return syn::Error::new_spanned(&input, "ToDataFrame only supports structs")
.to_compile_error()
.into();
}
};
let mut column_names: Vec<String> = Vec::new();
let mut column_values: Vec<TokenStream2> = Vec::new();
for field in fields.iter() {
let field_name = field.ident.as_ref().unwrap();
let field_name_str = to_snake_case(&field_name.to_string());
let field_type = &field.ty;
if let Some(value_expr) = generate_column_value(field_name, field_type) {
column_names.push(field_name_str);
column_values.push(value_expr);
}
}
let mut vec_column_values: Vec<TokenStream2> = Vec::new();
for field in fields.iter() {
let field_name = field.ident.as_ref().unwrap();
let field_type = &field.ty;
if let Some(value_expr) = generate_vec_column_value(field_name, field_type) {
vec_column_values.push(value_expr);
}
}
let expanded = quote! {
#[cfg(feature = "dataframe")]
impl #name {
pub fn to_dataframe(&self) -> ::polars::prelude::PolarsResult<::polars::prelude::DataFrame> {
use ::polars::prelude::*;
df![
#( #column_names => #column_values ),*
]
}
pub fn vec_to_dataframe(items: &[Self]) -> ::polars::prelude::PolarsResult<::polars::prelude::DataFrame> {
use ::polars::prelude::*;
df![
#( #column_names => #vec_column_values ),*
]
}
}
};
TokenStream::from(expanded)
}
fn to_snake_case(s: &str) -> String {
s.to_string()
}
fn generate_column_value(field_name: &syn::Ident, field_type: &Type) -> Option<TokenStream2> {
match field_type {
Type::Path(type_path) if is_string(type_path) => {
Some(quote! { [self.#field_name.as_str()] })
}
Type::Path(type_path) if is_formatted_value(type_path) => {
Some(quote! { [self.#field_name.raw] })
}
Type::Path(type_path) if is_option(type_path) => {
let inner_type = get_option_inner_type(type_path)?;
generate_option_value(field_name, inner_type)
}
Type::Path(type_path) if is_primitive(type_path) => Some(quote! { [self.#field_name] }),
_ => None,
}
}
fn generate_vec_column_value(field_name: &syn::Ident, field_type: &Type) -> Option<TokenStream2> {
match field_type {
Type::Path(type_path) if is_string(type_path) => {
Some(quote! { items.iter().map(|item| item.#field_name.as_str()).collect::<Vec<_>>() })
}
Type::Path(type_path) if is_formatted_value(type_path) => {
Some(quote! { items.iter().map(|item| item.#field_name.raw).collect::<Vec<_>>() })
}
Type::Path(type_path) if is_option(type_path) => {
let inner_type = get_option_inner_type(type_path)?;
generate_vec_option_value(field_name, inner_type)
}
Type::Path(type_path) if is_primitive(type_path) => {
Some(quote! { items.iter().map(|item| item.#field_name).collect::<Vec<_>>() })
}
_ => None,
}
}
fn generate_vec_option_value(field_name: &syn::Ident, inner_type: &Type) -> Option<TokenStream2> {
match inner_type {
Type::Path(type_path) if is_string(type_path) => Some(
quote! { items.iter().map(|item| item.#field_name.as_deref()).collect::<Vec<_>>() },
),
Type::Path(type_path) if is_formatted_value(type_path) => Some(
quote! { items.iter().map(|item| item.#field_name.as_ref().and_then(|v| v.raw)).collect::<Vec<_>>() },
),
Type::Path(type_path) if is_primitive(type_path) => {
Some(quote! { items.iter().map(|item| item.#field_name).collect::<Vec<_>>() })
}
_ => None,
}
}
fn generate_option_value(field_name: &syn::Ident, inner_type: &Type) -> Option<TokenStream2> {
match inner_type {
Type::Path(type_path) if is_string(type_path) => {
Some(quote! { [self.#field_name.as_deref()] })
}
Type::Path(type_path) if is_formatted_value(type_path) => {
Some(quote! { [self.#field_name.as_ref().and_then(|v| v.raw)] })
}
Type::Path(type_path) if is_primitive(type_path) => Some(quote! { [self.#field_name] }),
_ => None,
}
}
fn is_string(type_path: &TypePath) -> bool {
type_path
.path
.segments
.last()
.map(|seg| seg.ident == "String")
.unwrap_or(false)
}
fn is_option(type_path: &TypePath) -> bool {
type_path
.path
.segments
.last()
.map(|seg| seg.ident == "Option")
.unwrap_or(false)
}
fn is_formatted_value(type_path: &TypePath) -> bool {
type_path
.path
.segments
.last()
.map(|seg| seg.ident == "FormattedValue")
.unwrap_or(false)
}
fn is_primitive(type_path: &TypePath) -> bool {
type_path
.path
.segments
.last()
.map(|seg| {
let name = seg.ident.to_string();
matches!(
name.as_str(),
"i32" | "i64" | "f64" | "bool" | "u32" | "u64"
)
})
.unwrap_or(false)
}
fn get_option_inner_type(type_path: &TypePath) -> Option<&Type> {
let segment = type_path.path.segments.last()?;
if segment.ident != "Option" {
return None;
}
match &segment.arguments {
PathArguments::AngleBracketed(args) => args.args.first().and_then(|arg| {
if let GenericArgument::Type(ty) = arg {
Some(ty)
} else {
None
}
}),
_ => None,
}
}