finance_query_derive/

lib.rs

//! Derive macros for finance-query library.
//!
//! Provides `#[derive(ToDataFrame)]` for automatic DataFrame conversion.

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,
};

/// Derive macro for automatic DataFrame conversion.
///
/// Generates a `to_dataframe(&self) -> PolarsResult<DataFrame>` method
/// that converts all struct fields to DataFrame columns.
///
/// # Supported Types
///
/// - `String` → String column
/// - `Option<String>` → nullable String column
/// - `Option<FormattedValue<f64>>` → extracts `.raw` as `Option<f64>`
/// - `Option<FormattedValue<i64>>` → extracts `.raw` as `Option<i64>`
/// - `i32`, `i64`, `f64`, `bool` → direct columns
/// - `Option<T>` for primitives → nullable columns
/// - Nested structs/Vec → skipped (complex types not suitable for flat DataFrame)
///
/// # Example
///
/// ```ignore
/// #[derive(ToDataFrame)]
/// pub struct Quote {
///     pub symbol: String,
///     pub price: Option<FormattedValue<f64>>,
/// }
///
/// // Generates:
/// impl Quote {
///     pub fn to_dataframe(&self) -> PolarsResult<DataFrame> {
///         df![
///             "symbol" => [self.symbol.as_str()],
///             "price" => [self.price.as_ref().and_then(|v| v.raw)],
///         ]
///     }
/// }
/// ```
#[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);
        }
        // Skip fields that return None (complex nested types)
    }

    // Generate vec column value expressions (for vec_to_dataframe)
    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 {
            /// Converts this struct to a single-row polars DataFrame.
            ///
            /// All scalar fields are included as columns. Nested objects
            /// and complex types are excluded.
            ///
            /// This method is auto-generated by the `ToDataFrame` derive macro.
            pub fn to_dataframe(&self) -> ::polars::prelude::PolarsResult<::polars::prelude::DataFrame> {
                use ::polars::prelude::*;
                df![
                    #( #column_names => #column_values ),*
                ]
            }

            /// Converts a slice of structs to a multi-row polars DataFrame.
            ///
            /// All scalar fields are included as columns. Nested objects
            /// and complex types are excluded.
            ///
            /// This method is auto-generated by the `ToDataFrame` derive macro.
            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)
}

/// Converts a field name to snake_case for DataFrame column names.
fn to_snake_case(s: &str) -> String {
    s.to_string()
}

/// Generates the value expression for a DataFrame column based on field type.
///
/// Returns `None` for complex types that should be skipped.
fn generate_column_value(field_name: &syn::Ident, field_type: &Type) -> Option<TokenStream2> {
    match field_type {
        // Direct String
        Type::Path(type_path) if is_string(type_path) => {
            Some(quote! { [self.#field_name.as_str()] })
        }

        // Direct FormattedValue<T> - extract .raw
        Type::Path(type_path) if is_formatted_value(type_path) => {
            Some(quote! { [self.#field_name.raw] })
        }

        // Option<T>
        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)
        }

        // Direct primitives: i32, i64, f64, bool
        Type::Path(type_path) if is_primitive(type_path) => Some(quote! { [self.#field_name] }),

        // Skip all other types (Vec, nested structs, etc.)
        _ => None,
    }
}

/// Generates the value expression for a DataFrame column when iterating over a Vec.
///
/// Returns `None` for complex types that should be skipped.
fn generate_vec_column_value(field_name: &syn::Ident, field_type: &Type) -> Option<TokenStream2> {
    match field_type {
        // Direct String
        Type::Path(type_path) if is_string(type_path) => {
            Some(quote! { items.iter().map(|item| item.#field_name.as_str()).collect::<Vec<_>>() })
        }

        // Direct FormattedValue<T> - extract .raw
        Type::Path(type_path) if is_formatted_value(type_path) => {
            Some(quote! { items.iter().map(|item| item.#field_name.raw).collect::<Vec<_>>() })
        }

        // Option<T>
        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)
        }

        // Direct primitives: i32, i64, f64, bool
        Type::Path(type_path) if is_primitive(type_path) => {
            Some(quote! { items.iter().map(|item| item.#field_name).collect::<Vec<_>>() })
        }

        // Skip all other types (Vec, nested structs, etc.)
        _ => None,
    }
}

/// Generates value expression for Option<T> fields when iterating over a Vec.
fn generate_vec_option_value(field_name: &syn::Ident, inner_type: &Type) -> Option<TokenStream2> {
    match inner_type {
        // Option<String>
        Type::Path(type_path) if is_string(type_path) => Some(
            quote! { items.iter().map(|item| item.#field_name.as_deref()).collect::<Vec<_>>() },
        ),

        // Option<FormattedValue<T>> - extract .raw
        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<_>>() },
        ),

        // Option<primitive>
        Type::Path(type_path) if is_primitive(type_path) => {
            Some(quote! { items.iter().map(|item| item.#field_name).collect::<Vec<_>>() })
        }

        // Skip complex Option<T> types
        _ => None,
    }
}

/// Generates value expression for Option<T> fields.
fn generate_option_value(field_name: &syn::Ident, inner_type: &Type) -> Option<TokenStream2> {
    match inner_type {
        // Option<String>
        Type::Path(type_path) if is_string(type_path) => {
            Some(quote! { [self.#field_name.as_deref()] })
        }

        // Option<FormattedValue<T>> - extract .raw
        Type::Path(type_path) if is_formatted_value(type_path) => {
            Some(quote! { [self.#field_name.as_ref().and_then(|v| v.raw)] })
        }

        // Option<primitive>
        Type::Path(type_path) if is_primitive(type_path) => Some(quote! { [self.#field_name] }),

        // Skip complex Option<T> types
        _ => None,
    }
}

/// Checks if a type path is `String`.
fn is_string(type_path: &TypePath) -> bool {
    type_path
        .path
        .segments
        .last()
        .map(|seg| seg.ident == "String")
        .unwrap_or(false)
}

/// Checks if a type path is `Option<T>`.
fn is_option(type_path: &TypePath) -> bool {
    type_path
        .path
        .segments
        .last()
        .map(|seg| seg.ident == "Option")
        .unwrap_or(false)
}

/// Checks if a type path is `FormattedValue<T>`.
fn is_formatted_value(type_path: &TypePath) -> bool {
    type_path
        .path
        .segments
        .last()
        .map(|seg| seg.ident == "FormattedValue")
        .unwrap_or(false)
}

/// Checks if a type path is a primitive type (i32, i64, f64, bool).
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)
}

/// Extracts the inner type from Option<T>.
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,
    }
}