parfor/ome-metadata-derive/src/lib.rs

use proc_macro::TokenStream;
use quote::quote;
use syn::{Data, DeriveInput, Fields, Visibility, parse_macro_input};

#[proc_macro_derive(OmeXML)]
pub fn ome_metadata_derive_xml(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let ident = input.ident;

    let expanded = quote! {
        impl #ident {
            pub fn from_xml<T: AsRef<str>>(s: T) -> Result<Self, Error> {
                Ok(from_str(s.as_ref())?)
            }

            pub fn to_xml(&self) -> Result<String, Error> {
                Ok(to_string(self)?)
            }
        }
    };

    TokenStream::from(expanded)
}

#[proc_macro_derive(PyOmeComplexEnum)]
pub fn ome_metadata_derive_complex_enum(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let ident = input.ident;

    let mut variants = Vec::new();

    match input.data {
        Data::Enum(s) => {
            for variant in &s.variants {
                let variant_ident = &variant.ident;
                variants.push(quote! {
                    Self::#variant_ident(v) => v.clone().into_bound_py_any(py)?,
                });
            }
        }
        _ => {
            return syn::Error::new_spanned(ident, "PyOmeComplexEnum only supports structs")
                .to_compile_error()
                .into();
        }
    };

    let expanded = quote! {
        impl PyDisplay for #ident {
            fn py_str(&self) -> String {
                format!("{:?}", self)
            }
        }

        #[pyo3::prelude::pymethods]
        impl #ident {
            /// all possible variants of this enum that can be constructed or converted into
            #[staticmethod]
            fn variants() -> Vec<String> {
                Self::VARIANTS.iter().map(|v| v.trim_matches('"').to_string()).collect()
            }

            fn __repr__(&self) -> String {
                self.py_str()
            }

            fn __str__(&self) -> String {
                self.py_str()
            }

            fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>,)> {
                Ok((match self { #(#variants)* },))
            }

            #[staticmethod]
            #[pyo3(name = "from_xml")]
            fn py_from_xml(xml: &str) -> PyResult<Self> {
                Ok(Self::from_xml(xml)?)
            }

            #[pyo3(name = "to_xml")]
            fn py_to_xml(&self) -> PyResult<String> {
                Ok(self.to_xml()?)
            }
        }
    };

    TokenStream::from(expanded)
}

#[proc_macro_derive(PyOmeEnum)]
pub fn ome_metadata_derive_enum(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = input.ident;

    let expanded = quote! {
        impl PyDisplay for #name {
            fn py_str(&self) -> String {
                format!("{}", self)
            }
        }

        #[pyo3::prelude::pymethods]
        impl #name {
            #[new]
            fn new(unit: &str) -> PyResult<Self> {
                unit.parse().map_err(|_| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("could not parse {}", unit)))
            }

            /// all possible variants of this enum that can be constructed or converted into
            #[staticmethod]
            fn variants() -> Vec<String> {
                Self::VARIANTS.iter().map(|v| v.trim_matches('"').to_string()).collect()
            }

            fn __repr__(&self) -> String {
                self.py_str()
            }

            fn __str__(&self) -> String {
                self.py_str()
            }

            fn __getnewargs__(&self) -> (String,) {
                (format!("{}", self),)
            }

            #[staticmethod]
            #[pyo3(name = "from_xml")]
            fn py_from_xml(xml: &str) -> PyResult<Self> {
                Ok(Self::from_xml(xml)?)
            }

            #[pyo3(name = "to_xml")]
            fn py_to_xml(&self) -> PyResult<String> {
                Ok(self.to_xml()?)
            }
        }
    };

    TokenStream::from(expanded)
}

#[proc_macro_derive(PyOmeUnit)]
pub fn ome_metadata_derive_unit(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = input.ident;

    let expanded = quote! {
        impl PyDisplay for #name {
            fn py_str(&self) -> String {
                format!("{}", self)
            }
        }

        #[pyo3::prelude::pymethods]
        impl #name {
            #[new]
            fn new(unit: &str) -> PyResult<Self> {
                unit.parse().map_err(|_| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("could not parse {}", unit)))
            }

            /// convert a value between units
            #[pyo3(name = "convert")]
            fn py_convert(&self, unit: &Self, value: f64) -> PyResult<f64> {
                Ok(self.convert(unit, value)?)
            }

            /// all possible variants of this enum that can be constructed or converted into
            #[staticmethod]
            fn variants() -> Vec<String> {
                Self::VARIANTS.iter().map(|v| v.trim_matches('"').to_string()).collect()
            }

            fn __repr__(&self) -> String {
                self.py_str()
            }

            fn __str__(&self) -> String {
                self.py_str()
            }

            fn __getnewargs__(&self) -> (String,) {
                (format!("{}", self),)
            }

            #[staticmethod]
            #[pyo3(name = "from_xml")]
            fn py_from_xml(xml: &str) -> PyResult<Self> {
                Ok(Self::from_xml(xml)?)
            }

            #[pyo3(name = "to_xml")]
            fn py_to_xml(&self) -> PyResult<String> {
                Ok(self.to_xml()?)
            }
        }
    };

    TokenStream::from(expanded)
}

#[proc_macro_derive(PyOme)]
pub fn ome_metadata_derive(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = input.ident;

    // Only support named structs
    let fields = match input.data {
        Data::Struct(s) => match s.fields {
            Fields::Named(named) => named.named,
            _ => {
                return syn::Error::new_spanned(name, "PyOme only supports named structs")
                    .to_compile_error()
                    .into();
            }
        },
        _ => {
            return syn::Error::new_spanned(name, "PyOme only supports structs")
                .to_compile_error()
                .into();
        }
    };

    let mut idents = Vec::new();
    let mut types = Vec::new();
    let mut f = Vec::new();
    let mut ps = Vec::new();
    let mut new_kwargs = Vec::new();
    let mut s = Vec::new();

    for field in fields.iter() {
        let ty = &field.ty;
        if let Some(ident) = &field.ident
            && let Visibility::Public(_) = field.vis
        {
            let tmp = ident.to_string();
            let field_name = tmp.trim_start_matches("r#");
            idents.push(quote! {
                self.#ident
            });
            types.push(quote! { #ty });
            f.push(quote! {
                pub #ident: #ty
            });
            ps.push(quote! {
                self.#ident.py_str()
            });
            new_kwargs.push(quote! {
                (#field_name, self.#ident.clone().into_bound_py_any(py)?)
            });
            s.push(format!("{}: {{}}", field_name));
        } else {
            return syn::Error::new_spanned(field, "PyOme type only supports public named fields")
                .to_compile_error()
                .into();
        }
    }
    let fmt = format!("{}({})", name, s.join(", "));

    // Generate the token stream
    let expanded = quote! {
        impl PyDisplay for #name {
            fn py_str(&self) -> String {
                format!(#fmt, #(#ps,)*)
            }
        }

        #[pyo3::prelude::pymethods]
        impl #name {
            fn __getnewargs_ex__<'py>(&self, py: Python<'py>) -> PyResult<((), std::collections::HashMap<&str, Bound<'py, PyAny>>)> {
                Ok(((), std::collections::HashMap::from_iter([#(#new_kwargs,)*])))
            }

            fn __repr__(&self) -> String {
                self.py_str()
            }

            fn __str__(&self) -> String {
                self.py_str()
            }

            #[staticmethod]
            #[pyo3(name = "from_xml")]
            fn py_from_xml(xml: &str) -> PyResult<Self> {
                Ok(Self::from_xml(xml)?)
            }

            #[pyo3(name = "to_xml")]
            fn py_to_xml(&self) -> PyResult<String> {
                Ok(self.to_xml()?)
            }
        }
    };

    TokenStream::from(expanded)
}