Derive macro for complex enums (#208)

Add new attribute macro `gen_stub_pyclass_complex_enum` that creates
internal classes for the variants of rich enums, that is, enums with
variants that have fields.

Also creates in type_info.rs PyComplexEnumInfo and VariantInfo for
manual creation of stub information for rich enums.

Both of these are handled by 'ClassDef', which can now have internal
classes and will write properly indented internal classes to the stub
file.

This PR also fixes the implementation on tuples so that tuples of length
one are supported.

---------

Co-authored-by: Erik von Reis <[email protected]>
Co-authored-by: James Gregory <[email protected]>
Co-authored-by: Toshiki Teramura <[email protected]>

Author: 4 people

Cargo.lock

@@ -45,6 +45,100 @@ class B(A):
 class MyDate(datetime.date):
     ...
 
+class NumberComplex:
+    class FLOAT(NumberComplex):
+        r"""
+        Float variant
+        """
+        __match_args__ = ("_0",)
+        @property
+        def _0(self) -> builtins.float: ...
+        def __new__(cls, _0:builtins.float) -> NumberComplex.FLOAT: ...
+        def __len__(self) -> builtins.int: ...
+        def __getitem__(self, key:builtins.int) -> typing.Any: ...
+    
+    class INTEGER(NumberComplex):
+        r"""
+        Integer variant
+        """
+        __match_args__ = ("int",)
+        @property
+        def int(self) -> builtins.int:
+            r"""
+            The integer value
+            """
+        def __new__(cls, int:builtins.int=2) -> NumberComplex.INTEGER: ...
+    
+    ...
+
+class Shape1:
+    r"""
+    Example from PyO3 documentation for complex enum
+    https://pyo3.rs/v0.25.1/class.html#complex-enums
+    """
+    class Circle(Shape1):
+        __match_args__ = ("radius",)
+        @property
+        def radius(self) -> builtins.float: ...
+        def __new__(cls, radius:builtins.float) -> Shape1.Circle: ...
+    
+    class Rectangle(Shape1):
+        __match_args__ = ("width", "height",)
+        @property
+        def width(self) -> builtins.float: ...
+        @property
+        def height(self) -> builtins.float: ...
+        def __new__(cls, width:builtins.float, height:builtins.float) -> Shape1.Rectangle: ...
+    
+    class RegularPolygon(Shape1):
+        __match_args__ = ("_0", "_1",)
+        @property
+        def _0(self) -> builtins.int: ...
+        @property
+        def _1(self) -> builtins.float: ...
+        def __new__(cls, _0:builtins.int, _1:builtins.float) -> Shape1.RegularPolygon: ...
+        def __len__(self) -> builtins.int: ...
+        def __getitem__(self, key:builtins.int) -> typing.Any: ...
+    
+    class Nothing(Shape1):
+        __match_args__ = ((),)
+        def __new__(cls) -> Shape1.Nothing: ...
+    
+    ...
+
+class Shape2:
+    r"""
+    Example from PyO3 documentation for complex enum
+    https://pyo3.rs/v0.25.1/class.html#complex-enums
+    """
+    class Circle(Shape2):
+        __match_args__ = ("radius",)
+        @property
+        def radius(self) -> builtins.float: ...
+        def __new__(cls, radius:builtins.float=1.0) -> Shape2.Circle: ...
+    
+    class Rectangle(Shape2):
+        __match_args__ = ("width", "height",)
+        @property
+        def width(self) -> builtins.float: ...
+        @property
+        def height(self) -> builtins.float: ...
+        def __new__(cls, *, width:builtins.float, height:builtins.float) -> Shape2.Rectangle: ...
+    
+    class RegularPolygon(Shape2):
+        __match_args__ = ("side_count", "radius",)
+        @property
+        def side_count(self) -> builtins.int: ...
+        @property
+        def radius(self) -> builtins.float: ...
+        def __new__(cls, side_count:builtins.int, radius:builtins.float=1.0) -> Shape2.RegularPolygon: ...
+    
+    class Nothing(Shape2):
+        __match_args__ = ((),)
+        def __new__(cls) -> Shape2.Nothing: ...
+    
+    ...
+
 class Number(Enum):
     FLOAT = ...
     INTEGER = ...

examples/pure/pure.pyi

@@ -163,6 +163,54 @@ pub enum NumberRenameAll {
     Integer,
 }
 
+#[gen_stub_pyclass_complex_enum]
+#[pyclass]
+#[pyo3(rename_all = "UPPERCASE")]
+#[derive(Debug, Clone)]
+pub enum NumberComplex {
+    /// Float variant
+    Float(f64),
+    /// Integer variant
+    #[pyo3(constructor = (int=2))]
+    Integer {
+        /// The integer value
+        int: i32,
+    },
+}
+
+/// Example from PyO3 documentation for complex enum
+/// https://pyo3.rs/v0.25.1/class.html#complex-enums
+#[gen_stub_pyclass_complex_enum]
+#[pyclass]
+enum Shape1 {
+    Circle { radius: f64 },
+    Rectangle { width: f64, height: f64 },
+    RegularPolygon(u32, f64),
+    Nothing {},
+}
+
+/// Example from PyO3 documentation for complex enum
+/// https://pyo3.rs/v0.25.1/class.html#complex-enums
+#[gen_stub_pyclass_complex_enum]
+#[pyclass]
+enum Shape2 {
+    #[pyo3(constructor = (radius=1.0))]
+    Circle {
+        radius: f64,
+    },
+    #[pyo3(constructor = (*, width, height))]
+    Rectangle {
+        width: f64,
+        height: f64,
+    },
+    #[pyo3(constructor = (side_count, radius=1.0))]
+    RegularPolygon {
+        side_count: u32,
+        radius: f64,
+    },
+    Nothing {},
+}
+
 #[gen_stub_pymethods]
 #[pymethods]
 impl Number {
@@ -199,6 +247,9 @@ fn pure(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_class::<MyDate>()?;
     m.add_class::<Number>()?;
     m.add_class::<NumberRenameAll>()?;
+    m.add_class::<NumberComplex>()?;
+    m.add_class::<Shape1>()?;
+    m.add_class::<Shape2>()?;
     m.add_function(wrap_pyfunction!(sum, m)?)?;
     m.add_function(wrap_pyfunction!(create_dict, m)?)?;
     m.add_function(wrap_pyfunction!(read_dict, m)?)?;

examples/pure/src/lib.rs

@@ -1,4 +1,13 @@
-from pure import sum, create_dict, read_dict, echo_path, ahash_dict
+from pure import (
+    sum,
+    create_dict,
+    read_dict,
+    echo_path,
+    ahash_dict,
+    NumberComplex,
+    Shape1,
+    Shape2,
+)
 import pytest
 import pathlib
 
@@ -36,6 +45,67 @@ def test_read_dict():
     )
 
 
+def test_number_complex():
+    i = NumberComplex.INTEGER(1)
+    f = NumberComplex.FLOAT(1.5)
+    assert i.int == 1
+    assert f._0 == 1.5
+    assert len(f) == 1
+    i2 = NumberComplex.INTEGER()
+    assert i2.int == 2
+
+
+# Test code for complex enum case from PyO3 document
+# https://pyo3.rs/v0.25.1/class.html#complex-enums
+def test_complex_enum_shape1():
+    circle = Shape1.Circle(radius=10.0)
+    square = Shape1.RegularPolygon(4, 10.0)
+
+    assert isinstance(circle, Shape1)
+    assert isinstance(circle, Shape1.Circle)
+    assert circle.radius == 10.0
+
+    assert isinstance(square, Shape1)
+    assert isinstance(square, Shape1.RegularPolygon)
+    assert square[0] == 4  # Gets _0 field
+    assert square[1] == 10.0  # Gets _1 field
+
+    def count_vertices(cls, shape):
+        match shape:
+            case cls.Circle():
+                return 0
+            case cls.Rectangle():
+                return 4
+            case cls.RegularPolygon(n):
+                return n
+            case cls.Nothing():
+                return 0
+
+    assert count_vertices(Shape1, circle) == 0
+    assert count_vertices(Shape1, square) == 4
+
+
+# Test code for complex enum case from PyO3 document
+# https://pyo3.rs/v0.25.1/class.html#complex-enums
+def test_complex_enum_shape2():
+    circle = Shape2.Circle()
+    assert isinstance(circle, Shape2)
+    assert isinstance(circle, Shape2.Circle)
+    assert circle.radius == 1.0
+
+    square = Shape2.Rectangle(width=1, height=1)
+    assert isinstance(square, Shape2)
+    assert isinstance(square, Shape2.Rectangle)
+    assert square.width == 1
+    assert square.height == 1
+
+    hexagon = Shape2.RegularPolygon(6)
+    assert isinstance(hexagon, Shape2)
+    assert isinstance(hexagon, Shape2.RegularPolygon)
+    assert hexagon.side_count == 6
+    assert hexagon.radius == 1
+
+
 def test_path():
     out = echo_path(pathlib.Path("test"))
     assert out == pathlib.Path("test")

examples/pure/tests/test_python.py

@@ -72,19 +72,22 @@ mod attr;
 mod member;
 mod method;
 mod pyclass;
+mod pyclass_complex_enum;
 mod pyclass_enum;
 mod pyfunction;
 mod pymethods;
 mod renaming;
 mod signature;
 mod stub_type;
 mod util;
+mod variant;
 
 use arg::*;
 use attr::*;
 use member::*;
 use method::*;
 use pyclass::*;
+use pyclass_complex_enum::*;
 use pyclass_enum::*;
 use pyfunction::*;
 use pymethods::*;
@@ -123,6 +126,18 @@ pub fn pyclass_enum(item: TokenStream2) -> Result<TokenStream2> {
     })
 }
 
+pub fn pyclass_complex_enum(item: TokenStream2) -> Result<TokenStream2> {
+    let inner = PyComplexEnumInfo::try_from(parse2::<ItemEnum>(item.clone())?)?;
+    let derive_stub_type = StubType::from(&inner);
+    Ok(quote! {
+        #item
+        #derive_stub_type
+        pyo3_stub_gen::inventory::submit! {
+            #inner
+        }
+    })
+}
+
 pub fn pymethods(item: TokenStream2) -> Result<TokenStream2> {
     let mut item_impl = parse2::<ItemImpl>(item)?;
     let inner = PyMethodsInfo::try_from(item_impl.clone())?;

pyo3-stub-gen-derive/src/gen_stub.rs

@@ -58,6 +58,7 @@ pub enum Attr {
     Set,
     SetAll,
     Module(String),
+    Constructor(Signature),
     Signature(Signature),
     RenameAll(RenamingRule),
     Extends(Type),
@@ -148,6 +149,9 @@ pub fn parse_pyo3_attr(attr: &Attribute) -> Result<Vec<Attr>> {
                     [Ident(ident), Punct(_), Group(group)] => {
                         if ident == "signature" {
                             pyo3_attrs.push(Attr::Signature(syn::parse2(group.to_token_stream())?));
+                        } else if ident == "constructor" {
+                            pyo3_attrs
+                                .push(Attr::Constructor(syn::parse2(group.to_token_stream())?));
                         }
                     }
                     [Ident(ident), Punct(_), Ident(ident2)] => {

pyo3-stub-gen-derive/src/gen_stub/attr.rs

@@ -2,11 +2,12 @@ use crate::gen_stub::{attr::parse_gen_stub_default, extract_documents};
 
 use super::{escape_return_type, parse_pyo3_attrs, Attr};
 
+use crate::gen_stub::arg::ArgInfo;
 use proc_macro2::TokenStream as TokenStream2;
 use quote::{quote, ToTokens, TokenStreamExt};
 use syn::{Attribute, Error, Expr, Field, FnArg, ImplItemConst, ImplItemFn, Result, Type};
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct MemberInfo {
     doc: String,
     name: String,
@@ -196,3 +197,11 @@ impl ToTokens for MemberInfo {
         })
     }
 }
+
+impl From<MemberInfo> for ArgInfo {
+    fn from(value: MemberInfo) -> Self {
+        let MemberInfo { name, r#type, .. } = value;
+
+        Self { name, r#type }
+    }
+}

pyo3-stub-gen-derive/src/gen_stub/member.rs

@@ -1,9 +1,8 @@
+use super::{extract_documents, parse_pyo3_attrs, util::quote_option, Attr, MemberInfo, StubType};
 use proc_macro2::TokenStream as TokenStream2;
 use quote::{quote, ToTokens, TokenStreamExt};
 use syn::{parse_quote, Error, ItemStruct, Result, Type};
 
-use super::{extract_documents, parse_pyo3_attrs, util::quote_option, Attr, MemberInfo, StubType};
-
 pub struct PyClassInfo {
     pyclass_name: String,
     struct_type: Type,
@@ -142,7 +141,7 @@ mod test {
             "#,
         )?;
         let out = PyClassInfo::try_from(input)?.to_token_stream();
-        insta::assert_snapshot!(format_as_value(out), @r###"
+        insta::assert_snapshot!(format_as_value(out), @r#"
         ::pyo3_stub_gen::type_info::PyClassInfo {
             pyclass_name: "Placeholder",
             struct_id: std::any::TypeId::of::<PyPlaceholder>,
@@ -171,7 +170,7 @@ mod test {
             doc: "",
             bases: &[],
         }
-        "###);
+        "#);
         Ok(())
     }