diff --git a/tokenizers/display_derive/src/fmt_parsing.rs b/tokenizers/display_derive/src/fmt_parsing.rs index b9ebc2c06..006f63dcf 100644 --- a/tokenizers/display_derive/src/fmt_parsing.rs +++ b/tokenizers/display_derive/src/fmt_parsing.rs @@ -4,7 +4,7 @@ use quote::ToTokens; use syn::{ parse::{Parse, ParseStream}, punctuated::Punctuated, - token, Expr, + token, Attribute, Expr, }; /// Representation of a [`fmt`]-like attribute. @@ -18,7 +18,7 @@ pub struct FmtAttribute { /// Interpolation [`syn::LitStr`]. /// /// [`syn::LitStr`]: struct@syn::LitStr - lit: syn::LitStr, + pub lit: syn::LitStr, /// Optional [`token::Comma`]. /// @@ -26,7 +26,7 @@ pub struct FmtAttribute { comma: Option, /// Interpolation arguments. - args: Punctuated, + pub args: Punctuated, } impl Parse for FmtAttribute { @@ -67,11 +67,11 @@ impl ToTokens for FmtAttribute { /// in a [`FmtAttribute`]. /// This should be used in `[display(fmt="", alias=alias, expr)]`. /// [1]: https://doc.rust-lang.org/stable/std/fmt/index.html#named-parameters -struct FmtArgument { +pub struct FmtArgument { /// `identifier =` [`Ident`]. /// /// [`Ident`]: struct@syn::Ident - alias: Option<(syn::Ident, token::Eq)>, + pub alias: Option<(syn::Ident, token::Eq)>, /// `expression` [`Expr`]. expr: Expr, @@ -106,3 +106,20 @@ impl ToTokens for FmtArgument { self.expr.to_tokens(tokens) } } + +pub fn find_display_attribute(attrs: &[Attribute]) -> Option { + let display_attr = attrs.iter().find(|attr| attr.path.is_ident("display")); + + let attr: Option = if let Some(attr) = display_attr { + match attr.parse_args::() { + Ok(display_macro) => Some(display_macro), + Err(e) => { + e.to_compile_error(); + None + } + } + } else { + None + }; + attr +} diff --git a/tokenizers/display_derive/src/lib.rs b/tokenizers/display_derive/src/lib.rs index 80e0e538e..315135297 100644 --- a/tokenizers/display_derive/src/lib.rs +++ b/tokenizers/display_derive/src/lib.rs @@ -1,45 +1,31 @@ extern crate proc_macro; use proc_macro::TokenStream; use quote::quote; -use syn::{parse_macro_input, DeriveInput}; +use syn::{parse_macro_input, stringify_punct, DeriveInput}; mod fmt_parsing; -use fmt_parsing::FmtAttribute; +use fmt_parsing::{find_display_attribute, FmtAttribute}; #[proc_macro_derive(Display, attributes(display))] pub fn display_derive(input: TokenStream) -> TokenStream { // Parse the parsed_input tokens into a syntax tree let parsed_input = parse_macro_input!(input as DeriveInput); // Find the `display` attribute - let display_attr = parsed_input - .attrs - .iter() - .find(|attr| attr.path.is_ident("display")); - - let fmt = if let Some(attr) = display_attr { - match attr.parse_args::() { - Ok(display_macro) => quote! { write!(f, #display_macro) }, - Err(e) => return e.to_compile_error().into(), - } - } else { - quote! {} - }; + let attr = find_display_attribute(&parsed_input.attrs); // 1. If the attrs are not None, then we defer to this. // Meaning we juste return quote!{ format!(#fmt, #attr)} let ident = &parsed_input.ident; - let body = if fmt.is_empty() { + let body = { // 2. We automatically parse match &parsed_input.data { - syn::Data::Struct(s) => generate_fmt_impl_for_struct(s, ident), + syn::Data::Struct(s) => generate_fmt_impl_for_struct(s, ident, &attr), syn::Data::Enum(e) => generate_fmt_impl_for_enum(e, ident), syn::Data::Union(u) => { let error = syn::Error::new_spanned(u.union_token, "Unions are not supported"); return proc_macro::TokenStream::from(error.into_compile_error()); } } - } else { - fmt }; let expanded = quote! { @@ -57,41 +43,57 @@ pub fn display_derive(input: TokenStream) -> TokenStream { fn generate_fmt_impl_for_struct( data_struct: &syn::DataStruct, ident: &syn::Ident, + attrs: &Option, ) -> proc_macro2::TokenStream { let fields = &data_struct.fields; - // Extract field names and types - let field_names: Vec<_> = fields.iter().map(|f| f.ident.as_ref().unwrap()).collect(); - let field_types: Vec<_> = fields.iter().map(|f| &f.ty).collect(); - - quote! { - write!(f, "{}(", stringify!(#ident))?; - let mut first = true; - #( - if !first { - write!(f, ", ")?; - } - first = false; + // Generate field formatting expressions + let field_formats: Vec<_> = fields + .iter() + .map(|f| { + let field_name = &f.ident; + let fmts = find_display_attribute(&f.attrs); - let field_value = &self.#field_names; - write!(f, "{}=", stringify!(#field_names))?; - if std::any::TypeId::of::<#field_types>() == std::any::TypeId::of::() { - println!("We have a string!"); - write!(f, "\"{}\"", field_value)?; + if let Some(attr) = attrs { + if attr.args.is_empty() { + // If there is a prefix and no args, use fmts if it exists + if let Some(fmt) = fmts { + // Combine prefix and fmts + quote! { + write!(f, "{}{}", #fmt.lit.value(), #fmt.args.to_string())?; + } + } else { + // If no fmts, write just the field value + quote! { + write!(f, "{}", self.#field_name)?; + } + } + } else { + // If there are args to the attribute, use attr.lit and attr.args exclusively + quote! { + write!(f, "{}{}", #attr.lit.value(), #attr.args.to_string())?; + } + } } else { - let s = format!("{}", field_value); - let mut chars = s.chars(); - let mut prefix = (&mut chars).take(100 - 1).collect::(); - if chars.next().is_some() { - prefix.push('…'); + // If there is no attribute, print everything directly + quote! { + write!(f, "{}", self.#field_name)?; } - write!(f, "{}", prefix)?; } - )* - write!(f, ")") + }) + .collect(); + + // Generate the final implementation of Display trait for the struct + quote! { + impl std::fmt::Display for #ident { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{}(", stringify!(#ident))?; + #(#field_formats)* + write!(f, ")") + } + } } } - fn generate_fmt_impl_for_enum( data_enum: &syn::DataEnum, ident: &syn::Ident,