#![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, GenericParam, Ident, PathArguments, Type,
TypeParam, TypeParamBound, 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 has_format_param = input.generics.params.iter().any(|param| {
if let GenericParam::Type(TypeParam { bounds, .. }) = param {
bounds.iter().any(|b| {
if let TypeParamBound::Trait(tb) = b {
tb.path
.segments
.last()
.map(|s| s.ident == "Format")
.unwrap_or(false)
} else {
false
}
})
} else {
false
}
});
let impl_ty = if has_format_param {
quote! { #name<crate::format::Both> }
} else {
quote! { #name }
};
let format_param_ident: Option<&Ident> = input.generics.params.iter().find_map(|param| {
if let GenericParam::Type(TypeParam { ident, bounds, .. }) = param {
let is_format = bounds.iter().any(|b| {
if let TypeParamBound::Trait(tb) = b {
tb.path
.segments
.last()
.map(|s| s.ident == "Format")
.unwrap_or(false)
} else {
false
}
});
if is_format { Some(ident) } else { None }
} else {
None
}
});
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, format_param_ident)
{
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, format_param_ident)
{
vec_column_values.push(value_expr);
}
}
let expanded = quote! {
#[cfg(feature = "dataframe")]
impl #impl_ty {
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,
fmt_param: Option<&Ident>,
) -> 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, fmt_param)
}
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,
fmt_param: Option<&Ident>,
) -> 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, fmt_param)
}
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,
fmt_param: Option<&Ident>,
) -> 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) || is_format_assoc_value(type_path, fmt_param) =>
{
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,
fmt_param: Option<&Ident>,
) -> 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) || is_format_assoc_value(type_path, fmt_param) =>
{
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_format_assoc_value(type_path: &TypePath, fmt_param: Option<&Ident>) -> bool {
let Some(param) = fmt_param else { return false };
let segs = &type_path.path.segments;
segs.len() == 2 && segs[0].ident == *param && segs[1].ident == "Value"
}
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)
}
#[proc_macro_derive(FormatConvert)]
pub fn derive_format_convert(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,
"FormatConvert only supports structs with named fields",
)
.to_compile_error()
.into();
}
},
_ => {
return syn::Error::new_spanned(&input, "FormatConvert only supports structs")
.to_compile_error()
.into();
}
};
let format_param: Option<&Ident> = input.generics.params.iter().find_map(|param| {
if let GenericParam::Type(TypeParam { ident, bounds, .. }) = param {
let has_format = bounds.iter().any(|b| {
if let TypeParamBound::Trait(tb) = b {
tb.path
.segments
.last()
.map(|s| s.ident == "Format")
.unwrap_or(false)
} else {
false
}
});
if has_format { Some(ident) } else { None }
} else {
None
}
});
let Some(format_param) = format_param else {
return syn::Error::new_spanned(
&input,
"FormatConvert requires a generic param bounded by Format (e.g. <F: Format = Both>)",
)
.to_compile_error()
.into();
};
let classified: Vec<(&syn::Field, bool)> = fields
.iter()
.map(|f| (f, is_format_value_field(&f.ty, format_param)))
.collect();
let raw_field_exprs: Vec<_> = classified
.iter()
.map(|(f, is_fmt)| {
let ident = f.ident.as_ref().unwrap();
if *is_fmt {
quote! { #ident: v.#ident.and_then(|fv| fv.raw) }
} else {
quote! { #ident: v.#ident }
}
})
.collect();
let pretty_field_exprs: Vec<_> = classified
.iter()
.map(|(f, is_fmt)| {
let ident = f.ident.as_ref().unwrap();
if *is_fmt {
quote! { #ident: v.#ident.and_then(|fv| fv.fmt.or(fv.long_fmt)) }
} else {
quote! { #ident: v.#ident }
}
})
.collect();
let expanded = quote! {
impl From<#name<crate::format::Both>> for #name<crate::format::Raw> {
fn from(v: #name<crate::format::Both>) -> Self {
#name {
#(#raw_field_exprs,)*
}
}
}
impl From<#name<crate::format::Both>> for #name<crate::format::Pretty> {
fn from(v: #name<crate::format::Both>) -> Self {
#name {
#(#pretty_field_exprs,)*
}
}
}
impl #name<crate::format::Both> {
pub fn into_raw(self) -> #name<crate::format::Raw> { self.into() }
pub fn as_raw(&self) -> #name<crate::format::Raw> { self.clone().into() }
pub fn into_pretty(self) -> #name<crate::format::Pretty> { self.into() }
pub fn as_pretty(&self) -> #name<crate::format::Pretty> { self.clone().into() }
}
};
TokenStream::from(expanded)
}
fn is_format_value_field(ty: &Type, format_param: &Ident) -> bool {
let inner = match get_option_inner(ty) {
Some(t) => t,
None => return false,
};
if let Type::Path(tp) = inner {
let segs = &tp.path.segments;
if segs.len() == 2 {
return segs[0].ident == *format_param && segs[1].ident == "Value";
}
}
false
}
fn get_option_inner(ty: &Type) -> Option<&Type> {
if let Type::Path(tp) = ty {
let seg = tp.path.segments.last()?;
if seg.ident != "Option" {
return None;
}
if let PathArguments::AngleBracketed(args) = &seg.arguments {
return args.args.first().and_then(|a| {
if let GenericArgument::Type(t) = a {
Some(t)
} else {
None
}
});
}
}
None
}
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,
}
}