pymodule derive macro

derive/src/lib.rs

@@ -12,6 +12,7 @@ mod error;
 mod compile_bytecode;
 mod from_args;
 mod pyclass;
+mod pymodule;
 mod util;
 
 use error::{extract_spans, Diagnostic};
@@ -44,6 +45,13 @@ pub fn pyimpl(attr: TokenStream, item: TokenStream) -> TokenStream {
     result_to_tokens(pyclass::impl_pyimpl(attr, item))
 }
 
+#[proc_macro_attribute]
+pub fn pymodule(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let attr = parse_macro_input!(attr as AttributeArgs);
+    let item = parse_macro_input!(item as Item);
+    result_to_tokens(pymodule::impl_pymodule(attr, item))
+}
+
 #[proc_macro_attribute]
 pub fn pystruct_sequence(attr: TokenStream, item: TokenStream) -> TokenStream {
     let attr = parse_macro_input!(attr as AttributeArgs);

derive/src/pyclass.rs

@@ -1,5 +1,5 @@
 use super::Diagnostic;
-use crate::util::path_eq;
+use crate::util::{def_to_name, path_eq, strip_prefix, ItemMeta};
 use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::{quote, quote_spanned, ToTokens};
 use std::collections::{HashMap, HashSet};
@@ -42,181 +42,6 @@ enum ClassItem {
     },
 }
 
-struct ClassItemMeta<'a> {
-    sig: &'a Signature,
-    parent_type: &'static str,
-    meta: HashMap<String, Option<Lit>>,
-}
-
-impl<'a> ClassItemMeta<'a> {
-    const METHOD_NAMES: &'static [&'static str] = &["name", "magic"];
-    const PROPERTY_NAMES: &'static [&'static str] = &["name", "magic", "setter"];
-
-    fn from_nested_meta(
-        parent_type: &'static str,
-        sig: &'a Signature,
-        nested_meta: &[NestedMeta],
-        names: &[&'static str],
-    ) -> Result<Self, Diagnostic> {
-        let mut extracted = Self {
-            sig,
-            parent_type,
-            meta: HashMap::new(),
-        };
-
-        let validate_name = |name: &str, extracted: &Self| -> Result<(), Diagnostic> {
-            if names.contains(&name) {
-                if extracted.meta.contains_key(name) {
-                    bail_span!(
-                        &sig.ident,
-                        "#[{}] must have only one '{}'",
-                        parent_type,
-                        name
-                    );
-                } else {
-                    Ok(())
-                }
-            } else {
-                bail_span!(
-                    &sig.ident,
-                    "#[{}({})] is not one of allowed attributes {}",
-                    parent_type,
-                    name,
-                    names.join(", ")
-                );
-            }
-        };
-
-        for meta in nested_meta {
-            let meta = match meta {
-                NestedMeta::Meta(meta) => meta,
-                NestedMeta::Lit(_) => continue,
-            };
-
-            match meta {
-                Meta::NameValue(name_value) => {
-                    if let Some(ident) = name_value.path.get_ident() {
-                        let name = ident.to_string();
-                        validate_name(&name, &extracted)?;
-                        extracted.meta.insert(name, Some(name_value.lit.clone()));
-                    }
-                }
-                Meta::Path(path) => {
-                    if let Some(ident) = path.get_ident() {
-                        let name = ident.to_string();
-                        validate_name(&name, &extracted)?;
-                        extracted.meta.insert(name, None);
-                    } else {
-                        continue;
-                    }
-                }
-                _ => (),
-            }
-        }
-
-        Ok(extracted)
-    }
-
-    fn _str(&self, key: &str) -> Result<Option<String>, Diagnostic> {
-        Ok(match self.meta.get(key) {
-            Some(Some(lit)) => {
-                if let Lit::Str(s) = lit {
-                    Some(s.value())
-                } else {
-                    bail_span!(
-                        &self.sig.ident,
-                        "#[{}({} = ...)] must be a string",
-                        self.parent_type,
-                        key
-                    );
-                }
-            }
-            Some(None) => {
-                bail_span!(
-                    &self.sig.ident,
-                    "#[{}({} = ...)] is expected",
-                    self.parent_type,
-                    key,
-                );
-            }
-            None => None,
-        })
-    }
-
-    fn _bool(&self, key: &str) -> Result<bool, Diagnostic> {
-        Ok(match self.meta.get(key) {
-            Some(Some(_)) => {
-                bail_span!(
-                    &self.sig.ident,
-                    "#[{}({})] is expected",
-                    self.parent_type,
-                    key,
-                );
-            }
-            Some(None) => true,
-            None => false,
-        })
-    }
-
-    fn method_name(&self) -> Result<String, Diagnostic> {
-        let name = self._str("name")?;
-        let magic = self._bool("magic")?;
-        Ok(if let Some(name) = name {
-            name
-        } else {
-            let name = self.sig.ident.to_string();
-            if magic {
-                format!("__{}__", name)
-            } else {
-                name
-            }
-        })
-    }
-
-    fn setter(&self) -> Result<bool, Diagnostic> {
-        self._bool("setter")
-    }
-
-    fn property_name(&self) -> Result<String, Diagnostic> {
-        let magic = self._bool("magic")?;
-        let setter = self._bool("setter")?;
-        let name = self._str("name")?;
-
-        Ok(if let Some(name) = name {
-            name
-        } else {
-            let sig_name = self.sig.ident.to_string();
-            let name = if setter {
-                if let Some(name) = strip_prefix(&sig_name, "set_") {
-                    if name.is_empty() {
-                        bail_span!(
-                            &self.sig.ident,
-                            "A #[{}(setter)] fn with a set_* name must \
-                             have something after \"set_\"",
-                            self.parent_type
-                        )
-                    }
-                    name.to_string()
-                } else {
-                    bail_span!(
-                        &self.sig.ident,
-                        "A #[{}(setter)] fn must either have a `name` \
-                         parameter or a fn name along the lines of \"set_*\"",
-                        self.parent_type
-                    )
-                }
-            } else {
-                sig_name
-            };
-            if magic {
-                format!("__{}__", name)
-            } else {
-                name
-            }
-        })
-    }
-}
-
 impl Class {
     fn add_item(&mut self, item: ClassItem, span: Span) -> Result<(), Diagnostic> {
         if self.items.insert(item) {
@@ -238,12 +63,8 @@ impl Class {
             )
         })?;
 
-        let item_meta = ClassItemMeta::from_nested_meta(
-            "pymethod",
-            sig,
-            &nesteds,
-            ClassItemMeta::METHOD_NAMES,
-        )?;
+        let item_meta =
+            ItemMeta::from_nested_meta("pymethod", sig, &nesteds, ItemMeta::ATTRIBUTE_NAMES)?;
         Ok(ClassItem::Method {
             item_ident: sig.ident.clone(),
             py_name: item_meta.method_name()?,
@@ -258,12 +79,8 @@ impl Class {
                  #[pyclassmethod(name = \"...\")]",
             )
         })?;
-        let item_meta = ClassItemMeta::from_nested_meta(
-            "pyclassmethod",
-            sig,
-            &nesteds,
-            ClassItemMeta::METHOD_NAMES,
-        )?;
+        let item_meta =
+            ItemMeta::from_nested_meta("pyclassmethod", sig, &nesteds, ItemMeta::ATTRIBUTE_NAMES)?;
         Ok(ClassItem::ClassMethod {
             item_ident: sig.ident.clone(),
             py_name: item_meta.method_name()?,
@@ -278,12 +95,8 @@ impl Class {
                  #[pyproperty(name = \"...\")]"
             )
         })?;
-        let item_meta = ClassItemMeta::from_nested_meta(
-            "pyproperty",
-            sig,
-            &nesteds,
-            ClassItemMeta::PROPERTY_NAMES,
-        )?;
+        let item_meta =
+            ItemMeta::from_nested_meta("pyproperty", sig, &nesteds, ItemMeta::PROPERTY_NAMES)?;
         Ok(ClassItem::Property {
             py_name: item_meta.property_name()?,
             item_ident: sig.ident.clone(),
@@ -335,17 +148,16 @@ impl Class {
                 Some(name) => name,
                 None => continue,
             };
-            if name == "pymethod" {
-                self.add_item(Self::extract_method(sig, meta)?, meta_span)?;
-            } else if name == "pyclassmethod" {
-                self.add_item(Self::extract_classmethod(sig, meta)?, meta_span)?;
-            } else if name == "pyproperty" {
-                self.add_item(Self::extract_property(sig, meta)?, meta_span)?;
-            } else if name == "pyslot" {
-                self.add_item(Self::extract_slot(sig, meta)?, meta_span)?;
-            } else {
-                continue;
-            }
+            let item = match name.to_string().as_str() {
+                "pymethod" => Self::extract_method(sig, meta)?,
+                "pyclassmethod" => Self::extract_classmethod(sig, meta)?,
+                "pyproperty" => Self::extract_property(sig, meta)?,
+                "pyslot" => Self::extract_slot(sig, meta)?,
+                _ => {
+                    continue;
+                }
+            };
+            self.add_item(item, meta_span)?;
             attr_idxs.push(i);
         }
         let mut i = 0;
@@ -597,30 +409,9 @@ pub fn impl_pyimpl(attr: AttributeArgs, item: Item) -> Result<TokenStream2, Diag
 
 fn generate_class_def(
     ident: &Ident,
-    attr_name: &'static str,
-    attr: AttributeArgs,
+    name: &str,
     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 path_eq(&name_value.path, "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") {
@@ -646,7 +437,7 @@ fn generate_class_def(
 
     let ret = quote! {
         impl ::rustpython_vm::pyobject::PyClassDef for #ident {
-            const NAME: &'static str = #class_name;
+            const NAME: &'static str = #name;
             const DOC: Option<&'static str> = #doc;
         }
     };
@@ -663,7 +454,8 @@ pub fn impl_pyclass(attr: AttributeArgs, item: Item) -> Result<TokenStream2, Dia
         ),
     };
 
-    let class_def = generate_class_def(&ident, "pyclass", attr, &attrs)?;
+    let class_name = def_to_name(&ident, "pyclass", attr)?;
+    let class_def = generate_class_def(&ident, &class_name, &attrs)?;
 
     let ret = quote! {
         #item
@@ -681,7 +473,8 @@ pub fn impl_pystruct_sequence(attr: AttributeArgs, item: Item) -> Result<TokenSt
             "#[pystruct_sequence] can only be on a struct declaration"
         )
     };
-    let class_def = generate_class_def(&struc.ident, "pystruct_sequence", attr, &struc.attrs)?;
+    let class_name = def_to_name(&struc.ident, "pystruct_sequence", attr)?;
+    let class_def = generate_class_def(&struc.ident, &class_name, &struc.attrs)?;
     let mut properties = Vec::new();
     let mut field_names = Vec::new();
     for (i, field) in struc.fields.iter().enumerate() {
@@ -745,11 +538,3 @@ pub fn impl_pystruct_sequence(attr: AttributeArgs, item: Item) -> Result<TokenSt
     };
     Ok(ret)
 }
-
-fn strip_prefix<'a>(s: &'a str, prefix: &str) -> Option<&'a str> {
-    if s.starts_with(prefix) {
-        Some(&s[prefix.len()..])
-    } else {
-        None
-    }
-}

derive/src/pymodule.rs

@@ -0,0 +1,179 @@
+use super::Diagnostic;
+use crate::util::{def_to_name, ItemMeta};
+use proc_macro2::{Span, TokenStream as TokenStream2};
+use quote::{quote, quote_spanned};
+use std::collections::HashSet;
+use syn::{
+    parse_quote, spanned::Spanned, Attribute, AttributeArgs, Ident, Item, Meta, NestedMeta,
+    Signature,
+};
+
+fn meta_to_vec(meta: Meta) -> Result<Vec<NestedMeta>, Meta> {
+    match meta {
+        Meta::Path(_) => Ok(Vec::new()),
+        Meta::List(list) => Ok(list.nested.into_iter().collect()),
+        Meta::NameValue(_) => Err(meta),
+    }
+}
+
+#[derive(Default)]
+struct Module {
+    items: HashSet<ModuleItem>,
+}
+
+#[derive(PartialEq, Eq, Hash)]
+enum ModuleItem {
+    Function { item_ident: Ident, py_name: String },
+}
+
+impl Module {
+    fn add_item(&mut self, item: ModuleItem, span: Span) -> Result<(), Diagnostic> {
+        if self.items.insert(item) {
+            Ok(())
+        } else {
+            Err(Diagnostic::span_error(
+                span,
+                "Duplicate #[py*] attribute on pyimpl".to_owned(),
+            ))
+        }
+    }
+
+    fn extract_function(sig: &Signature, meta: Meta) -> Result<ModuleItem, Diagnostic> {
+        let nesteds = meta_to_vec(meta).map_err(|meta| {
+            err_span!(
+                meta,
+                "#[pyfunction = \"...\"] cannot be a name/value, you probably meant \
+                 #[pyfunction(name = \"...\")]",
+            )
+        })?;
+
+        let item_meta =
+            ItemMeta::from_nested_meta("pyfunction", sig, &nesteds, ItemMeta::SIMPLE_NAMES)?;
+        Ok(ModuleItem::Function {
+            item_ident: sig.ident.clone(),
+            py_name: item_meta.simple_name()?,
+        })
+    }
+
+    fn extract_item_from_syn(
+        &mut self,
+        attrs: &mut Vec<Attribute>,
+        sig: &Signature,
+    ) -> Result<(), Diagnostic> {
+        let mut attr_idxs = Vec::new();
+        for (i, meta) in attrs
+            .iter()
+            .filter_map(|attr| attr.parse_meta().ok())
+            .enumerate()
+        {
+            let meta_span = meta.span();
+            let name = match meta.path().get_ident() {
+                Some(name) => name,
+                None => continue,
+            };
+            if name == "pyfunction" {
+                self.add_item(Self::extract_function(sig, meta)?, meta_span)?;
+            } else {
+                continue;
+            }
+            attr_idxs.push(i);
+        }
+        let mut i = 0;
+        let mut attr_idxs = &*attr_idxs;
+        attrs.retain(|_| {
+            let drop = attr_idxs.first().copied() == Some(i);
+            if drop {
+                attr_idxs = &attr_idxs[1..];
+            }
+            i += 1;
+            !drop
+        });
+        for (i, idx) in attr_idxs.iter().enumerate() {
+            attrs.remove(idx - i);
+        }
+        Ok(())
+    }
+}
+
+struct ItemSig<'a> {
+    attrs: &'a mut Vec<Attribute>,
+    sig: &'a Signature,
+}
+
+fn extract_module_items(mut items: Vec<ItemSig>) -> Result<TokenStream2, Diagnostic> {
+    let mut diagnostics: Vec<Diagnostic> = Vec::new();
+
+    let mut class = Module::default();
+
+    for item in items.iter_mut() {
+        push_diag_result!(
+            diagnostics,
+            class.extract_item_from_syn(&mut item.attrs, item.sig),
+        );
+    }
+
+    let functions = class.items.into_iter().map(|item| match item {
+        ModuleItem::Function {
+            item_ident,
+            py_name,
+        } => {
+            let new_func = quote_spanned!(item_ident.span()=> .new_function(#item_ident));
+            quote! {
+                vm.__module_set_attr(&module, #py_name, vm.ctx#new_func).unwrap();
+            }
+        }
+    });
+
+    Diagnostic::from_vec(diagnostics)?;
+
+    Ok(quote! {
+        #(#functions)*
+    })
+}
+
+pub fn impl_pymodule(attr: AttributeArgs, item: Item) -> Result<TokenStream2, Diagnostic> {
+    match item {
+        Item::Mod(mut module) => {
+            let module_name = def_to_name(&module.ident, "pymodule", attr)?;
+
+            let content = &mut module.content.as_mut().unwrap().1;
+            let items = content
+                .iter_mut()
+                .filter_map(|item| match item {
+                    syn::Item::Fn(syn::ItemFn {
+                        attrs,
+                        vis: _vis,
+                        sig,
+                        ..
+                    }) => Some(ItemSig { attrs, sig }),
+                    _ => None,
+                })
+                .collect();
+
+            let extend_mod = extract_module_items(items)?;
+            content.push(parse_quote! {
+                pub(crate) fn extend_module(
+                    vm: &::rustpython_vm::vm::VirtualMachine,
+                    module: &::rustpython_vm::pyobject::PyObjectRef,
+                ) {
+                    #extend_mod
+                }
+            });
+            content.push(parse_quote! {
+                #[allow(dead_code)]
+                pub(crate) fn make_module(
+                    vm: &::rustpython_vm::vm::VirtualMachine
+                ) -> ::rustpython_vm::pyobject::PyObjectRef {
+                    let module = vm.new_module(#module_name, vm.ctx.new_dict());
+                    extend_module(vm, &module);
+                    module
+                }
+            });
+
+            Ok(quote! {
+                #module
+            })
+        }
+        other => bail_span!(other, "#[pymodule] can only be on a module declaration"),
+    }
+}

derive/src/util.rs

@@ -1,3 +1,223 @@
-pub fn path_eq(path: &syn::Path, s: &str) -> bool {
+use super::Diagnostic;
+use std::collections::HashMap;
+use syn::{AttributeArgs, Ident, Lit, Meta, NestedMeta, Path, Signature};
+
+pub fn path_eq(path: &Path, s: &str) -> bool {
     path.get_ident().map_or(false, |id| id == s)
 }
+
+pub fn def_to_name(
+    ident: &Ident,
+    attr_name: &'static str,
+    attr: AttributeArgs,
+) -> Result<String, Diagnostic> {
+    let mut name = None;
+    for attr in attr {
+        if let NestedMeta::Meta(meta) = attr {
+            if let Meta::NameValue(name_value) = meta {
+                if path_eq(&name_value.path, "name") {
+                    if let Lit::Str(s) = name_value.lit {
+                        name = Some(s.value());
+                    } else {
+                        bail_span!(
+                            name_value.lit,
+                            "#[{}(name = ...)] must be a string",
+                            attr_name
+                        );
+                    }
+                }
+            }
+        }
+    }
+    Ok(name.unwrap_or_else(|| ident.to_string()))
+}
+
+pub fn strip_prefix<'a>(s: &'a str, prefix: &str) -> Option<&'a str> {
+    if s.starts_with(prefix) {
+        Some(&s[prefix.len()..])
+    } else {
+        None
+    }
+}
+
+pub struct ItemMeta<'a> {
+    sig: &'a Signature,
+    parent_type: &'static str,
+    meta: HashMap<String, Option<Lit>>,
+}
+
+impl<'a> ItemMeta<'a> {
+    pub const SIMPLE_NAMES: &'static [&'static str] = &["name"];
+    pub const ATTRIBUTE_NAMES: &'static [&'static str] = &["name", "magic"];
+    pub const PROPERTY_NAMES: &'static [&'static str] = &["name", "magic", "setter"];
+
+    pub fn from_nested_meta(
+        parent_type: &'static str,
+        sig: &'a Signature,
+        nested_meta: &[NestedMeta],
+        names: &[&'static str],
+    ) -> Result<Self, Diagnostic> {
+        let mut extracted = Self {
+            sig,
+            parent_type,
+            meta: HashMap::new(),
+        };
+
+        let validate_name = |name: &str, extracted: &Self| -> Result<(), Diagnostic> {
+            if names.contains(&name) {
+                if extracted.meta.contains_key(name) {
+                    bail_span!(
+                        &sig.ident,
+                        "#[{}] must have only one '{}'",
+                        parent_type,
+                        name
+                    );
+                } else {
+                    Ok(())
+                }
+            } else {
+                bail_span!(
+                    &sig.ident,
+                    "#[{}({})] is not one of allowed attributes {}",
+                    parent_type,
+                    name,
+                    names.join(", ")
+                );
+            }
+        };
+
+        for meta in nested_meta {
+            let meta = match meta {
+                NestedMeta::Meta(meta) => meta,
+                NestedMeta::Lit(_) => continue,
+            };
+
+            match meta {
+                Meta::NameValue(name_value) => {
+                    if let Some(ident) = name_value.path.get_ident() {
+                        let name = ident.to_string();
+                        validate_name(&name, &extracted)?;
+                        extracted.meta.insert(name, Some(name_value.lit.clone()));
+                    }
+                }
+                Meta::Path(path) => {
+                    if let Some(ident) = path.get_ident() {
+                        let name = ident.to_string();
+                        validate_name(&name, &extracted)?;
+                        extracted.meta.insert(name, None);
+                    } else {
+                        continue;
+                    }
+                }
+                _ => (),
+            }
+        }
+
+        Ok(extracted)
+    }
+
+    fn _str(&self, key: &str) -> Result<Option<String>, Diagnostic> {
+        Ok(match self.meta.get(key) {
+            Some(Some(lit)) => {
+                if let Lit::Str(s) = lit {
+                    Some(s.value())
+                } else {
+                    bail_span!(
+                        &self.sig.ident,
+                        "#[{}({} = ...)] must be a string",
+                        self.parent_type,
+                        key
+                    );
+                }
+            }
+            Some(None) => {
+                bail_span!(
+                    &self.sig.ident,
+                    "#[{}({} = ...)] is expected",
+                    self.parent_type,
+                    key,
+                );
+            }
+            None => None,
+        })
+    }
+
+    fn _bool(&self, key: &str) -> Result<bool, Diagnostic> {
+        Ok(match self.meta.get(key) {
+            Some(Some(_)) => {
+                bail_span!(
+                    &self.sig.ident,
+                    "#[{}({})] is expected",
+                    self.parent_type,
+                    key,
+                );
+            }
+            Some(None) => true,
+            None => false,
+        })
+    }
+
+    pub fn simple_name(&self) -> Result<String, Diagnostic> {
+        Ok(self
+            ._str("name")?
+            .unwrap_or_else(|| self.sig.ident.to_string()))
+    }
+
+    pub fn method_name(&self) -> Result<String, Diagnostic> {
+        let name = self._str("name")?;
+        let magic = self._bool("magic")?;
+        Ok(if let Some(name) = name {
+            name
+        } else {
+            let name = self.sig.ident.to_string();
+            if magic {
+                format!("__{}__", name)
+            } else {
+                name
+            }
+        })
+    }
+
+    pub fn property_name(&self) -> Result<String, Diagnostic> {
+        let magic = self._bool("magic")?;
+        let setter = self._bool("setter")?;
+        let name = self._str("name")?;
+
+        Ok(if let Some(name) = name {
+            name
+        } else {
+            let sig_name = self.sig.ident.to_string();
+            let name = if setter {
+                if let Some(name) = strip_prefix(&sig_name, "set_") {
+                    if name.is_empty() {
+                        bail_span!(
+                            &self.sig.ident,
+                            "A #[{}(setter)] fn with a set_* name must \
+                             have something after \"set_\"",
+                            self.parent_type
+                        )
+                    }
+                    name.to_string()
+                } else {
+                    bail_span!(
+                        &self.sig.ident,
+                        "A #[{}(setter)] fn must either have a `name` \
+                         parameter or a fn name along the lines of \"set_*\"",
+                        self.parent_type
+                    )
+                }
+            } else {
+                sig_name
+            };
+            if magic {
+                format!("__{}__", name)
+            } else {
+                name
+            }
+        })
+    }
+
+    pub fn setter(&self) -> Result<bool, Diagnostic> {
+        self._bool("setter")
+    }
+}