flarrow_message_derive/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::{
6    DeriveInput, Field, Fields, Ident, Token, Variant, parse_macro_input, punctuated::Punctuated,
7    token::Comma,
8};
9
10#[proc_macro_derive(ArrowMessage)]
11pub fn from_into_arrow_derive(input: TokenStream) -> TokenStream {
12    let input = parse_macro_input!(input as DeriveInput);
13
14    let name = input.ident;
15
16    match input.data {
17        syn::Data::Struct(s) => match s.fields {
18            Fields::Named(ref fields) => struct_derive(name, fields.named.clone()),
19            _ => panic!("Only structs with named fields are supported"),
20        },
21        syn::Data::Enum(e) => enum_derive(name, e.variants.clone()),
22        _ => panic!("Only structs and enums are supported"),
23    }
24}
25
26fn struct_derive(name: Ident, fields: Punctuated<Field, Comma>) -> TokenStream {
27    let field_attributes = fields
28        .iter()
29        .map(|field| (&field.ident, &field.ty))
30        .collect::<Vec<_>>();
31
32    let fields_fill = field_attributes.iter().map(|&(field, ty)| {
33        quote! {
34            <#ty>::field(stringify!(#field)),
35        }
36    });
37
38    let union_data_fill = field_attributes.iter().map(|&(field, _)| {
39        quote! {
40            #field: extract_union_data(stringify!(#field), &map, &children)?,
41        }
42    });
43
44    let arrow_data_fill = field_attributes.iter().map(|&(field, _)| {
45        quote! {
46            self.#field.try_into_arrow()?,
47        }
48    });
49
50    let expanded = quote! {
51        impl ArrowMessage for #name {
52            fn field(name: impl Into<String>) -> arrow_schema::Field {
53                make_union_fields(
54                    name,
55                    vec![
56                        #(#fields_fill)*
57                    ],
58                )
59            }
60
61            fn try_from_arrow(data: arrow_data::ArrayData) -> ArrowResult<Self>
62            where
63                Self: Sized,
64            {
65                let (map, children) = unpack_union(data);
66
67                Ok(Self {
68                    #(#union_data_fill)*
69                })
70            }
71
72            fn try_into_arrow(self) -> ArrowResult<arrow_array::ArrayRef> {
73                let union_fields = get_union_fields::<Self>()?;
74
75                make_union_array(
76                    union_fields,
77                    vec![
78                        #(#arrow_data_fill)*
79                    ],
80                )
81            }
82        }
83
84        impl TryFrom<arrow_data::ArrayData> for #name {
85            type Error = arrow_schema::ArrowError;
86
87            fn try_from(data: arrow_data::ArrayData) -> ArrowResult<Self> {
88                #name::try_from_arrow(data)
89            }
90        }
91
92        impl TryFrom<#name> for arrow_data::ArrayData {
93            type Error = arrow_schema::ArrowError;
94
95            fn try_from(item: #name) -> ArrowResult<Self> {
96                item.try_into_arrow().map(|array| array.into_data())
97            }
98        }
99    };
100
101    TokenStream::from(expanded)
102}
103
104fn enum_derive(name: Ident, variants: Punctuated<Variant, Token![,]>) -> TokenStream {
105    let variants: Vec<_> = variants
106        .iter()
107        .map(|variant| {
108            let variant_name = &variant.ident;
109            let variant_str = variant_name.to_string().to_lowercase(); // Exemple : `Foo` -> "foo"
110            (variant_name, variant_str)
111        })
112        .collect();
113
114    let into_string_arms = variants.iter().map(|(variant_name, variant_str)| {
115        quote! {
116            #name::#variant_name => #variant_str.to_string(),
117        }
118    });
119
120    let try_from_string_arms = variants.iter().map(|(variant_name, variant_str)| {
121        quote! {
122            #variant_str => Ok(#name::#variant_name),
123        }
124    });
125
126    let expanded = quote! {
127        impl #name {
128            pub fn into_string(self) -> String {
129                match self {
130                    #(#into_string_arms)*
131                }
132            }
133
134            pub fn try_from_string(s: String) -> ArrowResult<Self> {
135                match s.as_str() {
136                    #(#try_from_string_arms)*
137                    _ => Err(arrow_schema::ArrowError::ParseError(format!("Invalid value for {}: {}", stringify!(#name), s))),
138                }
139            }
140        }
141
142        impl ArrowMessage for #name {
143            fn field(name: impl Into<String>) -> arrow_schema::Field {
144                String::field(name)
145            }
146
147            fn try_from_arrow(data: arrow_data::ArrayData) -> ArrowResult<Self>
148            where
149                Self: Sized,
150            {
151                Encoding::try_from_string(String::try_from_arrow(data)?)
152            }
153
154            fn try_into_arrow(self) -> ArrowResult<arrow_array::ArrayRef> {
155                String::try_into_arrow(self.into_string())
156            }
157        }
158
159
160        impl TryFrom<arrow_data::ArrayData> for #name {
161            type Error = arrow_schema::ArrowError;
162
163            fn try_from(data: arrow_data::ArrayData) -> ArrowResult<Self> {
164                #name::try_from_arrow(data)
165            }
166        }
167
168        impl TryFrom<#name> for arrow_data::ArrayData {
169            type Error = arrow_schema::ArrowError;
170
171            fn try_from(item: #name) -> ArrowResult<Self> {
172                item.try_into_arrow().map(|array| array.into_data())
173            }
174        }
175    };
176
177    TokenStream::from(expanded)
178}
flarrow_message_derive/lib.rs

flarrow_message_derive/
lib.rs
