RustPython/derive/src/pyclass.rs

use super::Diagnostic;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use std::collections::HashMap;
use syn::{
    Attribute, AttributeArgs, Ident, ImplItem, Index, Item, Lit, Meta, MethodSig, NestedMeta,
};

enum ClassItem {
    Method {
        item_ident: Ident,
        py_name: String,
    },
    ClassMethod {
        item_ident: Ident,
        py_name: String,
    },
    Property {
        item_ident: Ident,
        py_name: String,
        setter: bool,
    },
}

fn meta_to_vec(meta: Meta) -> Result<Vec<NestedMeta>, Meta> {
    match meta {
        Meta::Word(_) => Ok(Vec::new()),
        Meta::List(list) => Ok(list.nested.into_iter().collect()),
        Meta::NameValue(_) => Err(meta),
    }
}

impl ClassItem {
    fn extract_from_syn(
        attrs: &mut Vec<Attribute>,
        sig: &MethodSig,
    ) -> Result<Option<ClassItem>, Diagnostic> {
        let mut item = None;
        let mut attr_idx = None;
        for (i, meta) in attrs
            .iter()
            .filter_map(|attr| attr.parse_meta().ok())
            .enumerate()
        {
            let name = meta.name();
            if name == "pymethod" {
                if item.is_some() {
                    bail_span!(
                        sig.ident,
                        "You can only have one #[py*] attribute on an impl item"
                    )
                }
                let nesteds = meta_to_vec(meta).map_err(|meta| {
                    err_span!(
                        meta,
                        "#[pymethod = \"...\"] cannot be a name/value, you probably meant \
                         #[pymethod(name = \"...\")]",
                    )
                })?;
                let mut py_name = None;
                for meta in nesteds {
                    let meta = match meta {
                        NestedMeta::Meta(meta) => meta,
                        NestedMeta::Literal(_) => continue,
                    };
                    if let Meta::NameValue(name_value) = meta {
                        if name_value.ident == "name" {
                            if let Lit::Str(s) = &name_value.lit {
                                py_name = Some(s.value());
                            } else {
                                bail_span!(&sig.ident, "#[pymethod(name = ...)] must be a string");
                            }
                        }
                    }
                }
                item = Some(ClassItem::Method {
                    item_ident: sig.ident.clone(),
                    py_name: py_name.unwrap_or_else(|| sig.ident.to_string()),
                });
                attr_idx = Some(i);
            } else if name == "pyclassmethod" {
                if item.is_some() {
                    bail_span!(
                        sig.ident,
                        "You can only have one #[py*] attribute on an impl item"
                    )
                }
                let nesteds = meta_to_vec(meta).map_err(|meta| {
                    err_span!(
                        meta,
                        "#[pyclassmethod = \"...\"] cannot be a name/value, you probably meant \
                         #[pyclassmethod(name = \"...\")]",
                    )
                })?;
                let mut py_name = None;
                for meta in nesteds {
                    let meta = match meta {
                        NestedMeta::Meta(meta) => meta,
                        NestedMeta::Literal(_) => continue,
                    };
                    if let Meta::NameValue(name_value) = meta {
                        if name_value.ident == "name" {
                            if let Lit::Str(s) = &name_value.lit {
                                py_name = Some(s.value());
                            } else {
                                bail_span!(
                                    &sig.ident,
                                    "#[pyclassmethod(name = ...)] must be a string"
                                );
                            }
                        }
                    }
                }
                item = Some(ClassItem::ClassMethod {
                    item_ident: sig.ident.clone(),
                    py_name: py_name.unwrap_or_else(|| sig.ident.to_string()),
                });
                attr_idx = Some(i);
            } else if name == "pyproperty" {
                if item.is_some() {
                    bail_span!(
                        sig.ident,
                        "You can only have one #[py*] attribute on an impl item"
                    )
                }
                let nesteds = meta_to_vec(meta).map_err(|meta| {
                    err_span!(
                        meta,
                        "#[pyproperty = \"...\"] cannot be a name/value, you probably meant \
                         #[pyproperty(name = \"...\")]"
                    )
                })?;
                let mut setter = false;
                let mut py_name = None;
                for meta in nesteds {
                    let meta = match meta {
                        NestedMeta::Meta(meta) => meta,
                        NestedMeta::Literal(_) => continue,
                    };
                    match meta {
                        Meta::NameValue(name_value) => {
                            if name_value.ident == "name" {
                                if let Lit::Str(s) = &name_value.lit {
                                    py_name = Some(s.value());
                                } else {
                                    bail_span!(
                                        &sig.ident,
                                        "#[pyproperty(name = ...)] must be a string"
                                    );
                                }
                            }
                        }
                        Meta::Word(ident) => {
                            if ident == "setter" {
                                setter = true;
                            }
                        }
                        _ => {}
                    }
                }
                let py_name = match py_name {
                    Some(py_name) => py_name,
                    None => {
                        let item_ident = sig.ident.to_string();
                        if setter {
                            if item_ident.starts_with("set_") {
                                let name = &item_ident["set_".len()..];
                                if name.is_empty() {
                                    bail_span!(
                                        &sig.ident,
                                        "A #[pyproperty(setter)] fn with a set_* name must \
                                         have something after \"set_\""
                                    )
                                } else {
                                    name.to_string()
                                }
                            } else {
                                bail_span!(
                                    &sig.ident,
                                    "A #[pyproperty(setter)] fn must either have a `name` \
                                     parameter or a fn name along the lines of \"set_*\""
                                )
                            }
                        } else {
                            item_ident
                        }
                    }
                };
                item = Some(ClassItem::Property {
                    py_name,
                    item_ident: sig.ident.clone(),
                    setter,
                });
                attr_idx = Some(i);
            }
        }
        if let Some(attr_idx) = attr_idx {
            attrs.remove(attr_idx);
        }
        Ok(item)
    }
}

pub fn impl_pyimpl(_attr: AttributeArgs, item: Item) -> Result<TokenStream2, Diagnostic> {
    let mut imp = if let Item::Impl(imp) = item {
        imp
    } else {
        return Ok(quote!(#item));
    };

    let mut diagnostics: Vec<Diagnostic> = Vec::new();

    let items = imp
        .items
        .iter_mut()
        .filter_map(|item| {
            if let ImplItem::Method(meth) = item {
                ClassItem::extract_from_syn(&mut meth.attrs, &meth.sig)
                    .map_err(|err| diagnostics.push(err))
                    .unwrap_or_default()
            } else {
                None
            }
        })
        .collect::<Vec<_>>();
    let ty = &imp.self_ty;
    let mut properties: HashMap<&str, (Option<&Ident>, Option<&Ident>)> = HashMap::new();
    for item in items.iter() {
        if let ClassItem::Property {
            ref item_ident,
            ref py_name,
            setter,
        } = item
        {
            let entry = properties.entry(py_name).or_default();
            let func = if *setter { &mut entry.1 } else { &mut entry.0 };
            if func.is_some() {
                bail_span!(
                    item_ident,
                    "Multiple property accessors with name {:?}",
                    py_name
                )
            }
            *func = Some(item_ident);
        }
    }
    let methods = items.iter().filter_map(|item| match item {
        ClassItem::Method {
            item_ident,
            py_name,
        } => Some(quote! {
            class.set_str_attr(#py_name, ctx.new_rustfunc(Self::#item_ident));
        }),
        ClassItem::ClassMethod {
            item_ident,
            py_name,
        } => Some(quote! {
            class.set_str_attr(#py_name, ctx.new_classmethod(Self::#item_ident));
        }),
        _ => None,
    });
    let properties = properties
        .iter()
        .map(|(name, prop)| {
            let getter = match prop.0 {
                Some(getter) => getter,
                None => {
                    push_err_span!(
                        diagnostics,
                        prop.1.unwrap(),
                        "Property {:?} is missing a getter",
                        name
                    );
                    return TokenStream2::new();
                }
            };
            let add_setter = prop.1.map(|setter| quote!(.add_setter(Self::#setter)));
            quote! {
                class.set_str_attr(
                    #name,
                    ::rustpython_vm::obj::objproperty::PropertyBuilder::new(ctx)
                        .add_getter(Self::#getter)
                        #add_setter
                        .create(),
                );
            }
        })
        .collect::<Vec<_>>();

    Diagnostic::from_vec(diagnostics)?;

    let ret = quote! {
        #imp
        impl ::rustpython_vm::pyobject::PyClassImpl for #ty {
            fn impl_extend_class(
                ctx: &::rustpython_vm::pyobject::PyContext,
                class: &::rustpython_vm::obj::objtype::PyClassRef,
            ) {
                #(#methods)*
                #(#properties)*
            }
        }
    };
    Ok(ret)
}

fn generate_class_def(
    ident: &Ident,
    attr_name: &'static str,
    attr: AttributeArgs,
    attrs: &[Attribute],
) -> Result<TokenStream2, Diagnostic> {
    let mut class_name = None;
    for attr in attr {
        if let NestedMeta::Meta(meta) = attr {
            if let Meta::NameValue(name_value) = meta {
                if name_value.ident == "name" {
                    if let Lit::Str(s) = name_value.lit {
                        class_name = Some(s.value());
                    } else {
                        bail_span!(
                            name_value.lit,
                            "#[{}(name = ...)] must be a string",
                            attr_name
                        );
                    }
                }
            }
        }
    }
    let class_name = class_name.unwrap_or_else(|| ident.to_string());

    let mut doc: Option<Vec<String>> = None;
    for attr in attrs.iter() {
        if attr.path.is_ident("doc") {
            let meta = attr.parse_meta().expect("expected doc attr to be a meta");
            if let Meta::NameValue(name_value) = meta {
                if let Lit::Str(s) = name_value.lit {
                    let val = s.value().trim().to_string();
                    match doc {
                        Some(ref mut doc) => doc.push(val),
                        None => doc = Some(vec![val]),
                    }
                }
            }
        }
    }
    let doc = match doc {
        Some(doc) => {
            let doc = doc.join("\n");
            quote!(Some(#doc))
        }
        None => quote!(None),
    };

    let ret = quote! {
        impl ::rustpython_vm::pyobject::PyClassDef for #ident {
            const NAME: &'static str = #class_name;
            const DOC: Option<&'static str> = #doc;
        }
    };
    Ok(ret)
}

pub fn impl_pyclass(attr: AttributeArgs, item: Item) -> Result<TokenStream2, Diagnostic> {
    let (item, ident, attrs) = match item {
        Item::Struct(struc) => (quote!(#struc), struc.ident, struc.attrs),
        Item::Enum(enu) => (quote!(#enu), enu.ident, enu.attrs),
        other => bail_span!(
            other,
            "#[pyclass] can only be on a struct or enum declaration"
        ),
    };

    let class_def = generate_class_def(&ident, "pyclass", attr, &attrs)?;

    let ret = quote! {
        #item
        #class_def
    };
    Ok(ret)
}

pub fn impl_pystruct_sequence(attr: AttributeArgs, item: Item) -> Result<TokenStream2, Diagnostic> {
    let struc = if let Item::Struct(struc) = item {
        struc
    } else {
        bail_span!(
            item,
            "#[pystruct_sequence] can only be on a struct declaration"
        )
    };
    let class_def = generate_class_def(&struc.ident, "pystruct_sequence", attr, &struc.attrs)?;
    let mut properties = Vec::new();
    let mut field_names = Vec::new();
    for (i, field) in struc.fields.iter().enumerate() {
        let idx = Index::from(i);
        if let Some(ref field_name) = field.ident {
            let field_name_str = field_name.to_string();
            // TODO add doc to the generated property
            let property = quote! {
                class.set_str_attr(
                    #field_name_str,
                    ::rustpython_vm::obj::objproperty::PropertyBuilder::new(ctx)
                        .add_getter(|zelf: &::rustpython_vm::obj::objtuple::PyTuple,
                                     _vm: &::rustpython_vm::vm::VirtualMachine|
                                     zelf.fast_getitem(#idx))
                        .create(),
                );
            };
            properties.push(property);
            field_names.push(quote!(#field_name));
        } else {
            field_names.push(quote!(#idx));
        }
    }

    let ty = &struc.ident;
    let ret = quote! {
        #struc
        #class_def
        impl #ty {
            fn into_struct_sequence(&self,
                vm: &::rustpython_vm::vm::VirtualMachine,
                cls: ::rustpython_vm::obj::objtype::PyClassRef,
            ) -> ::rustpython_vm::pyobject::PyResult<::rustpython_vm::obj::objtuple::PyTupleRef> {
                let tuple = ::rustpython_vm::obj::objtuple::PyTuple::from(
                    vec![#(::rustpython_vm::pyobject::IntoPyObject::into_pyobject(
                        ::std::clone::Clone::clone(&self.#field_names),
                        vm,
                    )?),*],
                );
                ::rustpython_vm::pyobject::PyValue::into_ref_with_type(tuple, vm, cls)
            }
        }
        impl ::rustpython_vm::pyobject::PyClassImpl for #ty {
            fn impl_extend_class(
                ctx: &::rustpython_vm::pyobject::PyContext,
                class: &::rustpython_vm::obj::objtype::PyClassRef,
            ) {
                #(#properties)*
            }

            fn make_class(
                ctx: &::rustpython_vm::pyobject::PyContext
            ) -> ::rustpython_vm::obj::objtype::PyClassRef {
                let py_class = ctx.new_class(<Self as ::rustpython_vm::pyobject::PyClassDef>::NAME, ctx.tuple_type());
                Self::extend_class(ctx, &py_class);
                py_class
            }
        }
    };
    Ok(ret)
}