annotate_units.rs

use darling::Result;
use proc_macro2::TokenStream;
use quote::quote;
use syn::{parse_quote, DeriveInput, LitStr};

pub fn get_abbreviation(value: String) -> Option<&'static str> {
    // Any new value here must also be added in editoast_common/src/units.rs
    match value.replace("::option", "").as_str() {
        "meter" => Some("Length in m"),
        "millimeter" => Some("Length in mm"),
        "meter_per_second" => Some("Velocity in m·s⁻¹"),
        "kilometer_per_hour" => Some("Velocity in km·h⁻¹"),
        "meter_per_second_squared" => Some("Acceleration in m·s⁻²"),
        "kilogram" => Some("Mass in kg"),
        "newton" => Some("Force in N"),
        "hertz" => Some("Viscosity friction in N/(m·s⁻1) = s⁻¹"),
        "kilogram_per_meter" => Some("Aerodynamic drag in N/(m·s⁻1)² = kg·m⁻¹"),
        "per_meter" => Some("Aerodynamic drag per kg in (N/kg)/(m/s)² = m⁻¹"),
        _ => None,
    }
}

pub fn annotate_units(input: &mut DeriveInput) -> Result<TokenStream> {
    // We look for fields that have #[sered(with="meter")] attributes
    // and we push two new attributes to it to improve the OpenAPI
    if let syn::Data::Struct(s) = &mut input.data {
        for f in s.fields.iter_mut() {
            for attr in f.attrs.clone() {
                if attr.path().is_ident("serde") {
                    let _ = attr.parse_nested_meta(|meta| {
                        if meta.path.is_ident("with") {
                            let value = meta.value()?;
                            let s: LitStr = value.parse()?;
                            if let Some(abbreviation) = get_abbreviation(s.value()) {
                                if s.value().ends_with("::option") {
                                    f.attrs.push(
                                        parse_quote! {#[schema(value_type = Option<f64>)]},
                                    );
                                } else {
                                    f.attrs.push(parse_quote! {#[schema(value_type = f64)]});
                                }
                                f.attrs.push(parse_quote! {#[doc = #abbreviation]});
                            }
                        }
                        Ok(())
                    });
                }
            }
        }
    }
    Ok(quote! {#input})
}