Add support for pyo3(get(attr1, attr2, attr3))

Author: davidhozic

pyo3-macros-backend/src/attributes.rs

@@ -1,5 +1,6 @@
 use proc_macro2::TokenStream;
 use quote::{quote, ToTokens};
+use syn::parenthesized;
 use syn::parse::Parser;
 use syn::{
     ext::IdentExt,
@@ -339,6 +340,57 @@ impl ToTokens for NewImplTypeAttributeValue {
     }
 }
 
+#[derive(Clone)]
+pub struct GetListAttribute {
+    pub get_token: kw::get,
+    pub paren_token: syn::token::Paren,
+    pub fields: Punctuated<Ident, Token![,]>,
+}
+
+impl syn::parse::Parse for GetListAttribute {
+    fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
+        // Parse the keyword: get
+        let get_token: kw::get = input.parse()?;
+
+        // Parse the parentheses: ( ... )
+        let content;
+        let paren_token = parenthesized!(content in input);
+
+        // Parse identifiers inside: a, b, c
+        let fields =
+            content.parse_terminated(Ident::parse, Token![,])?;
+
+        // Reject empty list: get()
+        if fields.is_empty() {
+            return Err(syn::Error::new(
+                paren_token.span.join(),
+                "`get(...)` must contain at least one field name",
+            ));
+        }
+
+        Ok(GetListAttribute {
+            get_token,
+            paren_token,
+            fields,
+        })
+    }
+}
+
+impl ToTokens for GetListAttribute {
+    fn to_tokens(&self, tokens: &mut TokenStream) {
+        // keyword `get`
+        self.get_token.to_tokens(tokens);
+
+        // parentheses
+        let paren_content = self.fields.iter().map(|f| quote::quote! { #f });
+        let paren_tokens = quote::quote! { ( #( #paren_content ),* ) };
+
+        tokens.extend(quote::quote! {
+            #paren_tokens
+        });
+    }
+}
+
 pub type ExtendsAttribute = KeywordAttribute<kw::extends, Path>;
 pub type FreelistAttribute = KeywordAttribute<kw::freelist, Box<Expr>>;
 pub type ModuleAttribute = KeywordAttribute<kw::module, LitStr>;

pyo3-macros-backend/src/pyclass.rs

@@ -12,7 +12,7 @@ use crate::attributes::kw::frozen;
 use crate::attributes::{
     self, kw, take_pyo3_options, CrateAttribute, ExtendsAttribute, FreelistAttribute,
     ModuleAttribute, NameAttribute, NameLitStr, NewImplTypeAttribute, NewImplTypeAttributeValue,
-    RenameAllAttribute, StrFormatterAttribute,
+    RenameAllAttribute, StrFormatterAttribute, GetListAttribute
 };
 use crate::combine_errors::CombineErrors;
 #[cfg(feature = "experimental-inspect")]
@@ -97,6 +97,7 @@ pub struct PyClassPyO3Options {
     pub generic: Option<kw::generic>,
     pub from_py_object: Option<kw::from_py_object>,
     pub skip_from_py_object: Option<kw::skip_from_py_object>,
+    pub get: Option<GetListAttribute>
 }
 
 pub enum PyClassPyO3Option {
@@ -125,6 +126,7 @@ pub enum PyClassPyO3Option {
     Generic(kw::generic),
     FromPyObject(kw::from_py_object),
     SkipFromPyObject(kw::skip_from_py_object),
+    Get(GetListAttribute)
 }
 
 impl Parse for PyClassPyO3Option {
@@ -180,6 +182,8 @@ impl Parse for PyClassPyO3Option {
             input.parse().map(PyClassPyO3Option::FromPyObject)
         } else if lookahead.peek(attributes::kw::skip_from_py_object) {
             input.parse().map(PyClassPyO3Option::SkipFromPyObject)
+        } else if lookahead.peek(attributes::kw::get) {
+            input.parse().map(PyClassPyO3Option::Get)
         } else {
             Err(lookahead.error())
         }
@@ -274,6 +278,7 @@ impl PyClassPyO3Options {
                 );
                 set_option!(from_py_object)
             }
+            PyClassPyO3Option::Get(get) => set_option!(get)
         }
         Ok(())
     }
@@ -359,6 +364,24 @@ pub fn build_py_class(
         }
     }
 
+    if let Some(get_list_attr) = &args.options.get {
+        // get_list_attr contains the list of desired field names (NameAttribute or Ident)
+        for name in get_list_attr.fields.iter() {
+            // find matching field in `field_options`:
+            if let Some((_, field_opts)) =
+                field_options.iter_mut().find(|(f, _)| match &f.ident {
+                    Some(ident) => ident == name,
+                    None => false,
+                }) {
+                if let Some(old_get) = field_opts.get.replace(Annotated::Struct(kw::get_all::default())) {
+                    return Err(syn::Error::new(old_get.span(), "duplicate get specified"));
+                }
+            } else {
+                return Err(syn::Error::new_spanned(get_list_attr.clone(), format!("no field named `{}`", name)));
+            }
+        }
+    }
+
     impl_class(&class.ident, &args, doc, field_options, methods_type, ctx)
 }