humanbyte_derive/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::Span;
3use quote::quote;
4use syn::{parse_macro_input, DeriveInput};
5
6#[proc_macro_derive(HumanByte)]
7pub fn humanbyte(input: TokenStream) -> TokenStream {
8    let input_str = input.to_string();
9    let constructor = humanbyte_constructor(input_str.parse().unwrap());
10    let display = humanbyte_display(input_str.parse().unwrap());
11    let parse = humanbyte_parse(input_str.parse().unwrap());
12    let ops = humanbyte_ops(input_str.parse().unwrap());
13    let fromstr = humanbyte_fromstr(input_str.parse().unwrap());
14
15    let mut combined = format!("{}{}{}{}{}", constructor, display, parse, ops, fromstr);
16    if cfg!(feature = "serde") {
17        let serde = humanbyte_serde(input_str.parse().unwrap());
18        combined = format!("{}{}", combined, serde);
19    }
20    combined.parse().unwrap()
21}
22
23#[proc_macro_derive(HumanByteConstructor)]
24pub fn humanbyte_constructor(input: TokenStream) -> TokenStream {
25    let input = parse_macro_input!(input as DeriveInput);
26    let name = &input.ident;
27
28    // Define units with their multipliers and descriptions
29    let units = vec![
30        ("b", "1", "bytes"),
31        ("kb", "::humanbyte::KB", "kilobytes"),
32        ("kib", "::humanbyte::KIB", "kibibytes"),
33        ("mb", "::humanbyte::MB", "megabytes"),
34        ("mib", "::humanbyte::MIB", "mebibytes"),
35        ("gb", "::humanbyte::GB", "gigabytes"),
36        ("gib", "::humanbyte::GIB", "gibibytes"),
37        ("tb", "::humanbyte::TB", "terabytes"),
38        ("tib", "::humanbyte::TIB", "tebibytes"),
39        ("pb", "::humanbyte::PB", "petabytes"),
40        ("pib", "::humanbyte::PIB", "pebibytes"),
41    ];
42
43    // Generate methods
44    let methods = units.iter().map(|(fn_name, multiplier, description)| {
45        // Create an identifier for the method name
46        let method_name = syn::Ident::new(fn_name, Span::call_site());
47
48        // Parse the multiplier into an expression
49        let multiplier_expr: syn::Expr = syn::parse_str(multiplier).unwrap();
50
51        // Generate the documentation comment
52        let doc_comment = format!("Construct `{}` given an amount of {}.", name, description);
53
54        // Generate the method using quote!
55        quote! {
56            #[doc = #doc_comment]
57            #[inline(always)]
58            pub const fn #method_name(size: u64) -> Self {
59                Self(size * #multiplier_expr)
60            }
61        }
62    });
63
64    let expanded = quote! {
65        impl #name {
66            #(#methods)*
67        }
68
69        impl From<u64> for #name {
70            fn from(size: u64) -> #name {
71                Self(size)
72            }
73        }
74    };
75
76    TokenStream::from(expanded)
77}
78
79#[proc_macro_derive(HumanByteOps)]
80pub fn humanbyte_ops(input: TokenStream) -> TokenStream {
81    let input = parse_macro_input!(input as DeriveInput);
82    let name = &input.ident;
83
84    let expanded = quote! {
85        impl core::ops::Add<#name> for #name {
86            type Output = #name;
87
88            #[inline(always)]
89            fn add(self, rhs: #name) -> #name {
90                #name(self.0 + rhs.0)
91            }
92        }
93
94        impl core::ops::AddAssign<#name> for #name {
95            #[inline(always)]
96            fn add_assign(&mut self, rhs: #name) {
97                self.0 += rhs.0
98            }
99        }
100
101        impl<T> core::ops::Add<T> for #name
102        where
103            T: Into<u64>,
104        {
105            type Output = #name;
106            #[inline(always)]
107            fn add(self, rhs: T) -> #name {
108                #name(self.0 + (rhs.into()))
109            }
110        }
111
112        impl<T> core::ops::AddAssign<T> for #name
113        where
114            T: Into<u64>,
115        {
116            #[inline(always)]
117            fn add_assign(&mut self, rhs: T) {
118                self.0 += rhs.into();
119            }
120        }
121
122        impl core::ops::Sub<#name> for #name {
123            type Output = #name;
124
125            #[inline(always)]
126            fn sub(self, rhs: #name) -> #name {
127                #name(self.0 - rhs.0)
128            }
129        }
130
131        impl core::ops::SubAssign<#name> for #name {
132            #[inline(always)]
133            fn sub_assign(&mut self, rhs: #name) {
134                self.0 -= rhs.0
135            }
136        }
137
138        impl<T> core::ops::Sub<T> for #name
139        where
140            T: Into<u64>,
141        {
142            type Output = #name;
143
144            #[inline(always)]
145            fn sub(self, rhs: T) -> #name {
146                #name(self.0 - (rhs.into()))
147            }
148        }
149
150        impl<T> core::ops::SubAssign<T> for #name
151        where
152            T: Into<u64>,
153        {
154            #[inline(always)]
155            fn sub_assign(&mut self, rhs: T) {
156                self.0 -= rhs.into();
157            }
158        }
159
160        impl<T> core::ops::Mul<T> for #name
161        where
162            T: Into<u64>,
163        {
164            type Output = #name;
165            #[inline(always)]
166            fn mul(self, rhs: T) -> #name {
167                #name(self.0 * rhs.into())
168            }
169        }
170
171        impl<T> core::ops::MulAssign<T> for #name
172        where
173            T: Into<u64>,
174        {
175            #[inline(always)]
176            fn mul_assign(&mut self, rhs: T) {
177                self.0 *= rhs.into();
178            }
179        }
180
181        impl core::ops::Add<#name> for u64 {
182            type Output = #name;
183            #[inline(always)]
184            fn add(self, rhs: #name) -> #name {
185                #name(rhs.0 + self)
186            }
187        }
188
189        impl core::ops::Add<#name> for u32 {
190            type Output = #name;
191            #[inline(always)]
192            fn add(self, rhs: #name) -> #name {
193                #name(rhs.0 + (self as u64))
194            }
195        }
196
197        impl core::ops::Add<#name> for u16 {
198            type Output = #name;
199            #[inline(always)]
200            fn add(self, rhs: #name) -> #name {
201                #name(rhs.0 + (self as u64))
202            }
203        }
204
205        impl core::ops::Add<#name> for u8 {
206            type Output = #name;
207            #[inline(always)]
208            fn add(self, rhs: #name) -> #name {
209                #name(rhs.0 + (self as u64))
210            }
211        }
212
213        impl core::ops::Mul<#name> for u64 {
214            type Output = #name;
215            #[inline(always)]
216            fn mul(self, rhs: #name) -> #name {
217                #name(rhs.0 * self)
218            }
219        }
220
221        impl core::ops::Mul<#name> for u32 {
222            type Output = #name;
223            #[inline(always)]
224            fn mul(self, rhs: #name) -> #name {
225                #name(rhs.0 * (self as u64))
226            }
227        }
228
229        impl core::ops::Mul<#name> for u16 {
230            type Output = #name;
231            #[inline(always)]
232            fn mul(self, rhs: #name) -> #name {
233                #name(rhs.0 * (self as u64))
234            }
235        }
236
237        impl core::ops::Mul<#name> for u8 {
238            type Output = #name;
239            #[inline(always)]
240            fn mul(self, rhs: #name) -> #name {
241                #name(rhs.0 * (self as u64))
242            }
243        }
244
245        #[cfg(target_pointer_width = "64")]
246        impl core::ops::Add<#name> for usize {
247            type Output = #name;
248            #[inline(always)]
249            fn add(self, rhs: #name) -> #name {
250                #name(rhs.0 + (self as u64))
251            }
252        }
253
254
255        #[cfg(target_pointer_width = "64")]
256        impl core::ops::Sub<#name> for usize {
257            type Output = #name;
258            #[inline(always)]
259            fn sub(self, rhs: #name) -> #name {
260                #name(self as u64 - rhs.0)
261            }
262        }
263
264        #[cfg(target_pointer_width = "64")]
265        impl core::ops::Mul<#name> for usize {
266            type Output = #name;
267            #[inline(always)]
268            fn mul(self, rhs: #name) -> #name {
269                #name(rhs.0 * (self as u64))
270            }
271        }
272
273        impl #name {
274            /// Provides `HumanByteRange` with explicit lower and upper bounds.
275            pub fn range<I: Into<Self>>(start: I, stop: I) -> ::humanbyte::HumanByteRange<Self> {
276                ::humanbyte::HumanByteRange::new(Some(start), Some(stop))
277            }
278
279            /// Provides `HumanByteRange` with explicit lower bound. Upper bound is set to `u64::MAX`.
280            pub fn range_start<I: Into<Self>>(start: I) -> ::humanbyte::HumanByteRange<Self> {
281                ::humanbyte::HumanByteRange::new(Some(start), None)
282            }
283
284            /// Provides `HumanByteRange` with explicit lower bound. Upper bound is set to `u64::MAX`.
285            pub fn range_stop<I: Into<Self>>(stop: I) -> ::humanbyte::HumanByteRange<Self> {
286                ::humanbyte::HumanByteRange::new(None, Some(stop.into()))
287            }
288        }
289    };
290
291    TokenStream::from(expanded)
292}
293
294#[proc_macro_derive(HumanByteDisplay)]
295pub fn humanbyte_display(input: TokenStream) -> TokenStream {
296    let input = parse_macro_input!(input as DeriveInput);
297    let name = &input.ident;
298
299    let expanded = quote! {
300        impl core::fmt::Display for #name {
301            fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
302                f.pad(&::humanbyte::to_string(self.0, ::humanbyte::Format::IEC))
303            }
304        }
305
306        impl core::fmt::Debug for #name {
307            fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
308                write!(f, "{}", self)
309            }
310        }
311    };
312
313    TokenStream::from(expanded)
314}
315
316#[proc_macro_derive(HumanByteFromStr)]
317pub fn humanbyte_fromstr(input: TokenStream) -> TokenStream {
318    let input = parse_macro_input!(input as DeriveInput);
319    let name = &input.ident;
320
321    let expanded = quote! {
322        impl core::str::FromStr for #name {
323            type Err = ::humanbyte::String;
324
325            fn from_str(value: &str) -> core::result::Result<Self, Self::Err> {
326                if let Ok(v) = value.parse::<u64>() {
327                    return Ok(Self(v));
328                }
329                let number = ::humanbyte::take_while(value, |c| c.is_ascii_digit() || c == '.');
330                match number.parse::<f64>() {
331                    Ok(v) => {
332                        let suffix = ::humanbyte::skip_while(&value[number.len()..], char::is_whitespace);
333                        match suffix.parse::<::humanbyte::Unit>() {
334                            Ok(u) => Ok(Self((v * u64::from(u) as f64) as u64)),
335                            Err(error) => Err(::humanbyte::format!(
336                                "couldn't parse {:?} into a known SI unit, {}",
337                                suffix, error
338                            )),
339                        }
340                    }
341                    Err(error) => Err(::humanbyte::format!(
342                        "couldn't parse {:?} into a ByteSize, {}",
343                        value, error
344                    )),
345                }
346            }
347        }
348    };
349
350    TokenStream::from(expanded)
351}
352
353#[proc_macro_derive(HumanByteParse)]
354pub fn humanbyte_parse(input: TokenStream) -> TokenStream {
355    let input = parse_macro_input!(input as DeriveInput);
356    let name = &input.ident;
357
358    let expanded = quote! {
359        impl #name {
360            /// Returns the size as a string with an optional SI unit.
361            #[inline(always)]
362            pub fn to_string_as(&self, format: ::humanbyte::Format) -> ::humanbyte::String {
363                ::humanbyte::to_string(self.0, format)
364            }
365
366            /// Returns the inner u64 value.
367            #[inline(always)]
368            pub const fn as_u64(&self) -> u64 {
369                self.0
370            }
371
372            /// Returns the inner value as usize.
373            #[cfg(target_pointer_width = "64")]
374            #[inline(always)]
375            pub const fn as_usize(&self) -> usize {
376                self.0 as usize
377            }
378        }
379    };
380
381    TokenStream::from(expanded)
382}
383
384#[proc_macro_derive(HumanByteSerde)]
385pub fn humanbyte_serde(input: TokenStream) -> TokenStream {
386    let input = parse_macro_input!(input as DeriveInput);
387    let name = &input.ident;
388
389    let expanded = quote! {
390        impl<'de> ::humanbyte::serde::Deserialize<'de> for #name {
391            fn deserialize<D>(deserializer: D) -> core::result::Result<Self, D::Error>
392            where
393                D: ::humanbyte::serde::Deserializer<'de>,
394            {
395                struct ByteSizeVistor;
396
397                impl<'de> ::humanbyte::serde::de::Visitor<'de> for ByteSizeVistor {
398                    type Value = #name;
399
400                    fn expecting(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
401                        formatter.write_str("an integer or string")
402                    }
403
404                    fn visit_i64<E: ::humanbyte::serde::de::Error>(self, value: i64) -> core::result::Result<Self::Value, E> {
405                        if let Ok(val) = u64::try_from(value) {
406                            Ok(#name(val))
407                        } else {
408                            Err(E::invalid_value(
409                                ::humanbyte::serde::de::Unexpected::Signed(value),
410                                &"integer overflow",
411                            ))
412                        }
413                    }
414
415                    fn visit_u64<E: ::humanbyte::serde::de::Error>(self, value: u64) -> core::result::Result<Self::Value, E> {
416                        Ok(#name(value))
417                    }
418
419                    fn visit_str<E: ::humanbyte::serde::de::Error>(self, value: &str) -> core::result::Result<Self::Value, E> {
420                        if let Ok(val) = value.parse() {
421                            Ok(val)
422                        } else {
423                            Err(E::invalid_value(
424                                ::humanbyte::serde::de::Unexpected::Str(value),
425                                &"parsable string",
426                            ))
427                        }
428                    }
429                }
430
431                if deserializer.is_human_readable() {
432                    deserializer.deserialize_any(ByteSizeVistor)
433                } else {
434                    deserializer.deserialize_u64(ByteSizeVistor)
435                }
436            }
437        }
438        impl ::humanbyte::serde::Serialize for #name {
439            fn serialize<S>(&self, serializer: S) -> core::result::Result<S::Ok, S::Error>
440            where
441                S: ::humanbyte::serde::Serializer,
442            {
443                if serializer.is_human_readable() {
444                    <str>::serialize(self.to_string().as_str(), serializer)
445                } else {
446                    self.0.serialize(serializer)
447                }
448            }
449        }
450    };
451
452    TokenStream::from(expanded)
453}