Merge pull request #2804 from youknowone/static-metaclass

DimitrisJim · web-flow · commit a21ad9e390a0 · 2021-08-13T22:04:13.000+03:00
metaclass support for #[pyclass] macro
diff --git a/derive/src/pyclass.rs b/derive/src/pyclass.rs
@@ -142,6 +142,7 @@ fn generate_class_def(
     name: &str,
     module_name: Option<&str>,
     base: Option<String>,
+    metaclass: Option<String>,
     attrs: &[Attribute],
 ) -> std::result::Result<TokenStream, Diagnostic> {
     let doc = if let Some(doc) = attrs.doc() {
@@ -175,20 +176,31 @@ fn generate_class_def(
         )
         .into());
     }
-    let base = base.map(|name| Ident::new(&name, ident.span()));
-
     let base_class = if is_pystruct {
         quote! {
             fn static_baseclass() -> &'static ::rustpython_vm::builtins::PyTypeRef {
                 use rustpython_vm::StaticType;
                 rustpython_vm::builtins::PyTuple::static_type()
             }
         }
-    } else if let Some(base) = base {
+    } else if let Some(typ) = base {
+        let typ = Ident::new(&typ, ident.span());
         quote! {
             fn static_baseclass() -> &'static ::rustpython_vm::builtins::PyTypeRef {
                 use rustpython_vm::StaticType;
-                #base::static_type()
+                #typ::static_type()
+            }
+        }
+    } else {
+        quote!()
+    };
+
+    let meta_class = if let Some(typ) = metaclass {
+        let typ = Ident::new(&typ, ident.span());
+        quote! {
+            fn static_metaclass() -> &'static ::rustpython_vm::builtins::PyTypeRef {
+                use rustpython_vm::StaticType;
+                #typ::static_type()
             }
         }
     } else {
@@ -211,6 +223,8 @@ fn generate_class_def(
                 &CELL
             }
 
+            #meta_class
+
             #base_class
         }
     };
@@ -227,7 +241,15 @@ pub(crate) fn impl_pyclass(
     let class_name = class_meta.class_name()?;
     let module_name = class_meta.module()?;
     let base = class_meta.base()?;
-    let class_def = generate_class_def(ident, &class_name, module_name.as_deref(), base, attrs)?;
+    let metaclass = class_meta.metaclass()?;
+    let class_def = generate_class_def(
+        ident,
+        &class_name,
+        module_name.as_deref(),
+        base,
+        metaclass,
+        attrs,
+    )?;
 
     let ret = quote! {
         #item
diff --git a/derive/src/util.rs b/derive/src/util.rs
@@ -233,7 +233,7 @@ impl ItemMeta for SimpleItemMeta {
 pub(crate) struct ClassItemMeta(ItemMetaInner);
 
 impl ItemMeta for ClassItemMeta {
-    const ALLOWED_NAMES: &'static [&'static str] = &["module", "name", "base"];
+    const ALLOWED_NAMES: &'static [&'static str] = &["module", "name", "base", "metaclass"];
 
     fn from_inner(inner: ItemMetaInner) -> Self {
         Self(inner)
@@ -272,6 +272,10 @@ impl ClassItemMeta {
         self.inner()._optional_str("base")
     }
 
+    pub fn metaclass(&self) -> Result<Option<String>> {
+        self.inner()._optional_str("metaclass")
+    }
+
     pub fn module(&self) -> Result<Option<String>> {
         const KEY: &str = "module";
         let inner = self.inner();
