[Prototype] Python stubs generation

Cargo.toml

@@ -91,6 +91,9 @@ auto-initialize = []
 # Optimizes PyObject to Vec conversion and so on.
 nightly = []
 
+# Generates Python stubs (.pyi) files during compilation
+generate-stubs = ["pyo3-macros/generate-stubs"]
+
 # Activates all additional features
 # This is mostly intended for testing purposes - activating *all* of these isn't particularly useful.
 full = [

pyo3-macros-backend/Cargo.toml

@@ -25,3 +25,4 @@ features = ["derive", "parsing", "printing", "clone-impls", "full", "extra-trait
 [features]
 pyproto = []
 abi3 = []
+generate-stubs = []

pyo3-macros-backend/src/attributes.rs

@@ -26,6 +26,7 @@ pub mod kw {
     syn::custom_keyword!(signature);
     syn::custom_keyword!(subclass);
     syn::custom_keyword!(text_signature);
+    syn::custom_keyword!(type_signature);
     syn::custom_keyword!(transparent);
     syn::custom_keyword!(unsendable);
     syn::custom_keyword!(weakref);
@@ -82,6 +83,7 @@ pub type FreelistAttribute = KeywordAttribute<kw::freelist, Box<Expr>>;
 pub type ModuleAttribute = KeywordAttribute<kw::module, LitStr>;
 pub type NameAttribute = KeywordAttribute<kw::name, NameLitStr>;
 pub type TextSignatureAttribute = KeywordAttribute<kw::text_signature, LitStr>;
+pub type TypeSignatureAttribute = KeywordAttribute<kw::type_signature, LitStr>;
 
 impl<K: Parse + std::fmt::Debug, V: Parse> Parse for KeywordAttribute<K, V> {
     fn parse(input: ParseStream<'_>) -> Result<Self> {

pyo3-macros-backend/src/lib.rs

@@ -27,6 +27,7 @@ mod pymethod;
 #[cfg(feature = "pyproto")]
 mod pyproto;
 mod wrap;
+mod stubs;
 
 pub use frompyobject::build_derive_from_pyobject;
 pub use module::{process_functions_in_module, pymodule_impl, PyModuleOptions};

pyo3-macros-backend/src/pyclass.rs

@@ -1,20 +1,18 @@
 // Copyright (c) 2017-present PyO3 Project and Contributors
 
-use crate::attributes::{
-    self, kw, take_pyo3_options, CrateAttribute, ExtendsAttribute, FreelistAttribute,
-    ModuleAttribute, NameAttribute, NameLitStr, TextSignatureAttribute,
-};
+use crate::attributes::{self, kw, take_pyo3_options, CrateAttribute, ExtendsAttribute, FreelistAttribute, ModuleAttribute, NameAttribute, NameLitStr, TextSignatureAttribute, TypeSignatureAttribute};
 use crate::deprecations::{Deprecation, Deprecations};
 use crate::konst::{ConstAttributes, ConstSpec};
 use crate::pyimpl::{gen_default_items, gen_py_const, PyClassMethodsType};
 use crate::pymethod::{impl_py_getter_def, impl_py_setter_def, PropertyType};
 use crate::utils::{self, get_pyo3_crate, PythonDoc};
-use proc_macro2::{Span, TokenStream};
+use proc_macro2::{Span, TokenStream, TokenTree};
 use quote::quote;
 use syn::ext::IdentExt;
 use syn::parse::{Parse, ParseStream};
 use syn::punctuated::Punctuated;
 use syn::{parse_quote, spanned::Spanned, Result, Token};
+use syn::__private::str;
 
 /// If the class is derived from a Rust `struct` or `enum`.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
@@ -60,6 +58,7 @@ pub struct PyClassPyO3Options {
     pub name: Option<NameAttribute>,
     pub subclass: Option<kw::subclass>,
     pub text_signature: Option<TextSignatureAttribute>,
+    pub type_signature: Option<TypeSignatureAttribute>,
     pub unsendable: Option<kw::unsendable>,
     pub weakref: Option<kw::weakref>,
 
@@ -77,6 +76,7 @@ enum PyClassPyO3Option {
     Name(NameAttribute),
     Subclass(kw::subclass),
     TextSignature(TextSignatureAttribute),
+    TypeSignature(TypeSignatureAttribute),
     Unsendable(kw::unsendable),
     Weakref(kw::weakref),
 
@@ -106,6 +106,8 @@ impl Parse for PyClassPyO3Option {
             input.parse().map(PyClassPyO3Option::Subclass)
         } else if lookahead.peek(attributes::kw::text_signature) {
             input.parse().map(PyClassPyO3Option::TextSignature)
+        } else if lookahead.peek(attributes::kw::type_signature) {
+            input.parse().map(PyClassPyO3Option::TypeSignature)
         } else if lookahead.peek(attributes::kw::unsendable) {
             input.parse().map(PyClassPyO3Option::Unsendable)
         } else if lookahead.peek(attributes::kw::weakref) {
@@ -159,6 +161,7 @@ impl PyClassPyO3Options {
             PyClassPyO3Option::Name(name) => set_option!(name),
             PyClassPyO3Option::Subclass(subclass) => set_option!(subclass),
             PyClassPyO3Option::TextSignature(text_signature) => set_option!(text_signature),
+            PyClassPyO3Option::TypeSignature(type_signature) => set_option!(type_signature),
             PyClassPyO3Option::Unsendable(unsendable) => set_option!(unsendable),
             PyClassPyO3Option::Weakref(weakref) => set_option!(weakref),
 
@@ -213,6 +216,7 @@ pub fn build_py_class(
         }
     };
 
+    generate_stub_class(&class.ident, &args, &doc, &field_options, &methods_type, &krate);
     impl_class(&class.ident, &args, doc, field_options, methods_type, krate)
 }
 
@@ -221,12 +225,14 @@ struct FieldPyO3Options {
     get: bool,
     set: bool,
     name: Option<NameAttribute>,
+    type_signature: Option<TypeSignatureAttribute>
 }
 
 enum FieldPyO3Option {
     Get(attributes::kw::get),
     Set(attributes::kw::set),
     Name(NameAttribute),
+    TypeSignature(TypeSignatureAttribute),
 }
 
 impl Parse for FieldPyO3Option {
@@ -238,6 +244,8 @@ impl Parse for FieldPyO3Option {
             input.parse().map(FieldPyO3Option::Set)
         } else if lookahead.peek(attributes::kw::name) {
             input.parse().map(FieldPyO3Option::Name)
+        } else if lookahead.peek(attributes::kw::type_signature) {
+            input.parse().map(FieldPyO3Option::TypeSignature)
         } else {
             Err(lookahead.error())
         }
@@ -250,6 +258,7 @@ impl FieldPyO3Options {
             get: false,
             set: false,
             name: None,
+            type_signature: None,
         };
 
         for option in take_pyo3_options(attrs)? {
@@ -275,6 +284,13 @@ impl FieldPyO3Options {
                     );
                     options.name = Some(name);
                 }
+                FieldPyO3Option::TypeSignature(typing) => {
+                    ensure_spanned!(
+                        options.type_signature.is_none(),
+                        typing.span() => "`type_signature` may only be specified once"
+                    );
+                    options.type_signature = Some(typing);
+                }
             }
         }
 
@@ -971,3 +987,96 @@ fn define_inventory_class(inventory_class_name: &syn::Ident) -> TokenStream {
         _pyo3::inventory::collect!(#inventory_class_name);
     }
 }
+
+#[cfg_attr(not(feature = "generate-stubs"), allow(unused_variables))]
+fn generate_stub_class(
+    cls: &syn::Ident,
+    args: &PyClassArgs,
+    doc: &PythonDoc,
+    field_options: &Vec<(&syn::Field, FieldPyO3Options)>,
+    methods_type: &PyClassMethodsType,
+    krate: &syn::Path,
+) {
+    #[cfg(feature = "generate-stubs")] {
+        //TODO: declare inheritance, if required
+        println!("# ID: {}", cls);
+        println!("class {}:", get_class_python_name(cls, args).to_string());
+
+        let mut is_empty = true; // if the class is empty, it should contain 'pass'
+
+        // Generate the documentation
+        let doc: TokenStream = quote! {#doc};
+
+        let mut doc_iter = doc.into_iter();
+        {  // The first token should be 'concat'
+            let first_token = doc_iter.next()
+                .expect("All documentation comments start with a 'concat!' macro, but the documentation is empty");
+            match first_token {
+                TokenTree::Ident(ref ident) => assert_eq!("concat".to_string(), ident.to_string(), "All documentation comments start with a 'concat!' macro, but the first token is not 'concat': {:?}", first_token),
+                _ => panic!("All documentation comments start with a 'concat!' macro, but the first token is not an identifier: {:?}", first_token)
+            }
+        }
+        {   // The second token should be '!' (because it's the concat! macro)
+            let second_token = doc_iter.next()
+                .expect("All documentation comments start with a 'concat!' macro, but the documentation only has a single token");
+            match second_token {
+                TokenTree::Punct(ref punct) => assert_eq!('!', punct.as_char(), "All documentation comments start with a 'concat!' macro, but the word 'concat' is not followed by '!': {:?}", punct),
+                _ => panic!("All documentation comments start with a 'concat!' macro, but the second token is not a '!': {:?}", second_token)
+            }
+        }
+        {   // The third item should be a group that contains the actual documentation text
+            let third_token = doc_iter.next()
+                .expect("All documentation comments start with a 'concat!' macro, but it has no arguments");
+            match third_token {
+                TokenTree::Group(ref group) => {
+                    println!("\t\"\"\"");
+                    for line in group.stream() {
+                        if let TokenTree::Literal(ref literal) = line {
+                            let unescaped_literal = literal.to_string()
+                                .replace("\\'", "'");
+
+                            let line = unescaped_literal
+                                .trim_start_matches("\"")
+                                .trim_end_matches("\"");
+
+                            if line != "\\n" && line != "\\u{0}" {
+                                for sub_line in line.split("\\n") {
+                                    println!("\t{}", sub_line);
+                                }
+                            }
+                        }
+                    }
+                    println!("\t\"\"\"");
+                },
+                _ => panic!("All documentation comments start with a 'concat!' macro, but its arguments are not a group: {:?}", third_token)
+            }
+        }
+
+        // Generate the different top-level fields
+        for (field, options) in field_options {
+            if options.get || options.set {
+                is_empty = false;
+
+                let name = options.name.as_ref().map(|n| &n.value.0)
+                    .unwrap_or(field.ident.as_ref().expect("This field has neither the #[pyo3(name = ...)] macro applied to, nor a Rust name, it's impossible to create a stub for it"))
+                    .to_string();
+
+                print!("\t{}", name);
+
+                if let Some(typing) = &options.type_signature {
+                    print!(": {}", typing.value.value())
+                } else {
+                    print!(": Any")
+                }
+
+                println!();
+            }
+        }
+
+        //TODO: methods
+
+        if is_empty {
+            println!("\tpass");
+        }
+    }
+}

pyo3-macros-backend/src/pyfunction.rs

@@ -18,6 +18,7 @@ use syn::{
     parse::{Parse, ParseBuffer, ParseStream},
     token::Comma,
 };
+use crate::attributes::TypeSignatureAttribute;
 
 #[derive(Debug, Clone, PartialEq)]
 pub enum Argument {
@@ -238,6 +239,7 @@ pub struct PyFunctionOptions {
     pub name: Option<NameAttribute>,
     pub signature: Option<PyFunctionSignature>,
     pub text_signature: Option<TextSignatureAttribute>,
+    pub type_signature: Option<TypeSignatureAttribute>,
     pub deprecations: Deprecations,
     pub krate: Option<CrateAttribute>,
 }
@@ -278,6 +280,7 @@ pub enum PyFunctionOption {
     PassModule(attributes::kw::pass_module),
     Signature(PyFunctionSignature),
     TextSignature(TextSignatureAttribute),
+    TypeSignature(TypeSignatureAttribute),
     Crate(CrateAttribute),
 }
 
@@ -292,6 +295,8 @@ impl Parse for PyFunctionOption {
             input.parse().map(PyFunctionOption::Signature)
         } else if lookahead.peek(attributes::kw::text_signature) {
             input.parse().map(PyFunctionOption::TextSignature)
+        } else if lookahead.peek(attributes::kw::type_signature) {
+            input.parse().map(PyFunctionOption::TypeSignature)
         } else if lookahead.peek(syn::Token![crate]) {
             input.parse().map(PyFunctionOption::Crate)
         } else {
@@ -336,6 +341,13 @@ impl PyFunctionOptions {
                     );
                     self.text_signature = Some(text_signature);
                 }
+                PyFunctionOption::TypeSignature(type_signature) => {
+                    ensure_spanned!(
+                        self.type_signature.is_none(),
+                        type_signature.kw.span() => "`type_signature` may only be specified once"
+                    );
+                    self.type_signature = Some(type_signature);
+                }
                 PyFunctionOption::Crate(path) => {
                     ensure_spanned!(
                         self.krate.is_none(),

pyo3-macros-backend/src/pymethod.rs

@@ -16,6 +16,7 @@ use proc_macro2::{Span, TokenStream};
 use quote::{format_ident, quote, ToTokens};
 use syn::Ident;
 use syn::{ext::IdentExt, spanned::Spanned, Result};
+use crate::stubs::map_rust_type_to_python;
 
 pub enum GeneratedPyMethod {
     Method(TokenStream),
@@ -176,6 +177,8 @@ pub fn gen_py_method(
     let method = PyMethod::parse(sig, &mut *meth_attrs, options)?;
     let spec = &method.spec;
 
+    generate_stub_method(cls, spec);
+
     Ok(match (method.kind, &spec.tp) {
         // Class attributes go before protos so that class attributes can be used to set proto
         // method to None.
@@ -230,6 +233,52 @@ pub fn gen_py_method(
     })
 }
 
+#[cfg_attr(not(feature = "generate-stubs"), allow(unused_variables))]
+fn generate_stub_method(cls: &syn::Type, method: &FnSpec) {
+    #[cfg(feature = "generate-stubs")] {
+        println!("# ID: {}", cls.to_token_stream());
+
+        match &method.tp {
+            FnType::FnStatic => println!("\t@staticmethod"),
+            FnType::FnClass => println!("\t@classmethod"),
+            FnType::Getter(_) => println!("\t@property"),
+            FnType::Setter(_) => println!("\t@{}.setter", method.python_name),
+            _ => {},
+        }
+
+        print!("\tdef {}(", method.python_name);
+
+        let mut comma_required = false;
+        match &method.tp {
+            FnType::Fn(_) | FnType::Getter(_) | FnType::Setter(_) => {
+                comma_required = true;
+                print!("self")
+            },
+            FnType::FnClass | FnType::FnNew => {
+                comma_required = true;
+                print!("cls")
+            },
+            _ => {}
+        }
+
+        for arg in &method.args {
+            if comma_required {
+                print!(", ");
+            }
+
+            print!("{}: {}", arg.name, map_rust_type_to_python(&arg.ty));
+
+            if arg.optional.is_some() {
+                print!(" = ...");
+            }
+
+            comma_required = true;
+        }
+        println!(") -> {}: ...", map_rust_type_to_python(&method.output));
+
+    }
+}
+
 pub fn check_generic(sig: &syn::Signature) -> syn::Result<()> {
     let err_msg = |typ| format!("Python functions cannot have generic {} parameters", typ);
     for param in &sig.generics.params {