pyo3_special_method_derive 0.1.0

//! Derive macros to help with Rust PyO3 support.
//!
//! This crate automatically derives the following functions for structs and enums:
//! - `__str__`
//! - `__repr__`
//! - `__dir__`
//!
//! Note: The `StrReprHelper` macro requires `T: Debug` for each `T` inside the item.
//! The `Debug` trait is used for the outputs.
//!
//! - You can skip exposure of variants or fields with the `#[attr]` attribute
//! - You can skip variants or fields for `__str__` or `__repr__` differently with the `#[skip_str]` and `#[skip_repr]` attributes
//! - Struct fields which are not `pub` are skipped automatically
//!

extern crate proc_macro;
use dir::get_dir_enum_variants;
use proc_macro::TokenStream;
use quote::quote;
use str_repr::display_debug_derive;
use syn::{parse_macro_input, Data, DeriveInput, Fields, Visibility};

mod dir;
mod str_repr;

/// Add a `__dir__` method to the struct in a `#[pymethods]` impl.
///
/// - You can skip exposure of certain fields by adding the `#[skip]` attribute macro
/// - For structs, all fields are skipped which are not marked `pub`
///
/// ## Example
/// ```
/// use pyo3::pyclass;
/// use pyo3_special_method_derive::DirHelper;
/// #[pyclass]
/// #[derive(DirHelper)]
/// struct Person {
///     pub name: String,
///     address: String,
///     #[skip]
///     pub phone_number: String,
/// }
/// ```
#[proc_macro_derive(DirHelper, attributes(skip))]
pub fn dir_helper_derive(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    // Get the name of the struct
    let name = &input.ident;

    // Generate code to match the struct's fields
    let expanded = match input.data {
        Data::Struct(data) => {
            match data.fields {
                Fields::Named(fields) => {
                    // If the struct has named fields extract their names
                    let field_names = fields
                        .named
                        .iter()
                        .filter(|f| !f.attrs.iter().any(|attr| attr.path().is_ident("skip")))
                        .filter(|f| matches!(f.vis, Visibility::Public(_)))
                        .map(|f| f.ident.as_ref().unwrap())
                        .collect::<Vec<_>>();

                    if field_names.is_empty() {
                        quote! {
                            #[pyo3::pymethods]
                            impl #name {
                                pub fn __dir__(&self) -> Vec<String> {
                                    Vec::new()
                                }
                            }
                        }
                    } else {
                        // Prepare an array where the elements are expressions that prepare the field vec
                        let mut assigner = proc_macro2::TokenStream::new();
                        quote_into::quote_into!(assigner += [#{
                            for name in field_names {
                                quote_into::quote_into!(assigner += (names.push(stringify!(#name).to_string())),)
                            }
                        }];);
                        quote! {
                            #[pyo3::pymethods]
                            impl #name {
                                pub fn __dir__(&self) -> Vec<String> {
                                    let mut names = Vec::new();
                                    #assigner
                                    names
                                }
                            }
                        }
                    }
                }
                Fields::Unit => {
                    // If the struct has no fields
                    quote! {
                        #[pyo3::pymethods]
                        impl #name {
                            pub fn __dir__(&self) -> Vec<String> {
                                Vec::new()
                            }
                        }
                    }
                }
                Fields::Unnamed(_) => {
                    quote! {
                        compile_error!("Unnamed fields for struct are not supported for DirHelper derive.");
                    }
                }
            }
        }
        Data::Enum(e) => {
            let variants = get_dir_enum_variants(&e);
            let mut assigner = proc_macro2::TokenStream::new();
            quote_into::quote_into!(assigner += [#{
                for name in variants {
                    quote_into::quote_into!(assigner += (names.push(stringify!(#name).to_string())),)
                }
            }];);
            quote! {
                #[pyo3::pymethods]
                impl #name {
                    pub fn __dir__(&self) -> Vec<String> {
                        let mut names = Vec::new();
                        #assigner
                        names
                    }
                }
            }
        }
        Data::Union(_) => {
            quote! {
                compile_error!("Unions are not supported for DirHelper derive");
            }
        }
    };
    TokenStream::from(expanded)
}

/// Add `__str__` and `__repr__` methods to the struct in a `#[pymethods]` impl.
///
/// - You can skip printing of certain fields by adding the `#[skip]` attribute macro
/// - To specialze skipping depending on `__str__` and `__repr__`, you can use the `#[skip_str]`
/// and `#[skip_repr]` attributes which skip for `__str__` and `__repr__` respectively
/// - For structs, all fields are skipped which are not marked `pub`
///
/// ## Example
/// ```
/// use pyo3::pyclass;
/// use pyo3_special_method_derive::StrReprHelper;
/// #[pyclass]
/// #[derive(StrReprHelper)]
/// struct Person {
///     pub name: String,
///     address: String,
///     #[skip]
///     pub phone_number: String,
/// }
/// ```
#[proc_macro_derive(StrReprHelper, attributes(skip, skip_str, skip_repr))]
pub fn str_repr_helper_derive(input_stream: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input_stream as DeriveInput);

    // Get the name of the struct
    let name = &input.ident;

    let display_debug_derive_body = display_debug_derive(&input);

    let expanded = quote! {
        #display_debug_derive_body

        #[pyo3::pymethods]
        impl #name {
            pub fn __str__(&self) -> String {
                format!("{self}")
            }

            pub fn __repr__(&self) -> String {
                format!("{self:?}")
            }
        }
    };

    TokenStream::from(expanded)
}