Merge pull request #2034 from b05902132/derive-enum

#[pyclass] for fieldless enums

Author: davidhewitt

CHANGELOG.md

@@ -22,9 +22,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Add `Py::setattr` method. [#2009](https://github.com/PyO3/pyo3/pull/2009)
 - Add `PyCapsule`, exposing the [Capsule API](https://docs.python.org/3/c-api/capsule.html#capsules). [#1980](https://github.com/PyO3/pyo3/pull/1980)
-- All PyO3 proc-macros except the deprecated `#[pyproto]` now accept a supplemental attribute `#[pyo3(crate = "some::path")]` that specifies
-  where to find the `pyo3` crate, in case it has been renamed or is re-exported and not found at the crate root. [#2022](https://github.com/PyO3/pyo3/pull/2022)
 - Expose `pyo3-build-config` APIs for cross-compiling and Python configuration discovery for use in other projects. [#1996](https://github.com/PyO3/pyo3/pull/1996)
+- All PyO3 proc-macros except the deprecated `#[pyproto]` now accept a supplemental attribute `#[pyo3(crate = "some::path")]` that specifies where to find the `pyo3` crate, in case it has been renamed or is re-exported and not found at the crate root. [#2022](https://github.com/PyO3/pyo3/pull/2022)
+- Enable `#[pyclass]` for fieldless (aka C-like) enums. [#2034](https://github.com/PyO3/pyo3/pull/2034)
 - Add buffer magic methods `__getbuffer__` and `__releasebuffer__` to `#[pymethods]`. [#2067](https://github.com/PyO3/pyo3/pull/2067)
 - Accept paths in `wrap_pyfunction` and `wrap_pymodule`. [#2081](https://github.com/PyO3/pyo3/pull/2081)
 - Add check for correct number of arguments on magic methods. [#2083](https://github.com/PyO3/pyo3/pull/2083)

guide/src/class.md

@@ -2,7 +2,7 @@
 
 PyO3 exposes a group of attributes powered by Rust's proc macro system for defining Python classes as Rust structs.
 
-The main attribute is `#[pyclass]`, which is placed upon a Rust `struct` to generate a Python type for it. A struct will usually also have *one* `#[pymethods]`-annotated `impl` block for the struct, which is used to define Python methods and constants for the generated Python type. (If the [`multiple-pymethods`] feature is enabled each `#[pyclass]` is allowed to have multiple `#[pymethods]` blocks.) Finally, there may be multiple `#[pyproto]` trait implementations for the struct, which are used to define certain python magic methods such as `__str__`.
+The main attribute is `#[pyclass]`, which is placed upon a Rust `struct` or a fieldless `enum` (a.k.a. C-like enum) to generate a Python type for it. They will usually also have *one* `#[pymethods]`-annotated `impl` block for the struct, which is used to define Python methods and constants for the generated Python type. (If the [`multiple-pymethods`] feature is enabled each `#[pyclass]` is allowed to have multiple `#[pymethods]` blocks.) Finally, there may be multiple `#[pyproto]` trait implementations for the struct, which are used to define certain python magic methods such as `__str__`.
 
 This chapter will discuss the functionality and configuration these attributes offer. Below is a list of links to the relevant section of this chapter for each:
 
@@ -20,9 +20,7 @@ This chapter will discuss the functionality and configuration these attributes o
 
 ## Defining a new class
 
-To define a custom Python class, a Rust struct needs to be annotated with the
-`#[pyclass]` attribute.
-
+To define a custom Python class, add the `#[pyclass]` attribute to a Rust struct or a fieldless enum.
 ```rust
 # #![allow(dead_code)]
 # use pyo3::prelude::*;
@@ -31,11 +29,17 @@ struct MyClass {
     # #[pyo3(get)]
     num: i32,
 }
+
+#[pyclass]
+enum MyEnum {
+    Variant,
+    OtherVariant = 30, // PyO3 supports custom discriminants.
+}
 ```
 
-Because Python objects are freely shared between threads by the Python interpreter, all structs annotated with `#[pyclass]` must implement `Send` (unless annotated with [`#[pyclass(unsendable)]`](#customizing-the-class)).
+Because Python objects are freely shared between threads by the Python interpreter, all types annotated with `#[pyclass]` must implement `Send` (unless annotated with [`#[pyclass(unsendable)]`](#customizing-the-class)).
 
-The above example generates implementations for [`PyTypeInfo`], [`PyTypeObject`], and [`PyClass`] for `MyClass`. To see these generated implementations, refer to the [implementation details](#implementation-details) at the end of this chapter.
+The above example generates implementations for [`PyTypeInfo`], [`PyTypeObject`], and [`PyClass`] for `MyClass` and `MyEnum`. To see these generated implementations, refer to the [implementation details](#implementation-details) at the end of this chapter.
 
 ## Adding the class to a module
 
@@ -140,8 +144,8 @@ so that they can benefit from a freelist. `XXX` is a number of items for the fre
 * `gc` - Classes with the `gc` parameter participate in Python garbage collection.
 If a custom class contains references to other Python objects that can be collected, the [`PyGCProtocol`]({{#PYO3_DOCS_URL}}/pyo3/class/gc/trait.PyGCProtocol.html) trait has to be implemented.
 * `weakref` - Adds support for Python weak references.
-* `extends=BaseType` - Use a custom base class. The base `BaseType` must implement `PyTypeInfo`.
-* `subclass` - Allows Python classes to inherit from this class.
+* `extends=BaseType` - Use a custom base class. The base `BaseType` must implement `PyTypeInfo`. `enum` pyclasses can't use a custom base class.
+* `subclass` - Allows Python classes to inherit from this class. `enum` pyclasses can't be inherited from.
 * `dict` - Adds `__dict__` support, so that the instances of this type have a dictionary containing arbitrary instance variables.
 * `unsendable` - Making it safe to expose `!Send` structs to Python, where all object can be accessed
    by multiple threads. A class marked with `unsendable` panics when accessed by another thread.
@@ -351,7 +355,7 @@ impl SubClass {
 ## Object properties
 
 PyO3 supports two ways to add properties to your `#[pyclass]`:
-- For simple fields with no side effects, a `#[pyo3(get, set)]` attribute can be added directly to the field definition in the `#[pyclass]`.
+- For simple struct fields with no side effects, a `#[pyo3(get, set)]` attribute can be added directly to the field definition in the `#[pyclass]`.
 - For properties which require computation you can define `#[getter]` and `#[setter]` functions in the [`#[pymethods]`](#instance-methods) block.
 
 We'll cover each of these in the following sections.
@@ -802,6 +806,118 @@ impl MyClass {
 Note that `text_signature` on classes is not compatible with compilation in
 `abi3` mode until Python 3.10 or greater.
 
+## #[pyclass] enums
+
+Currently PyO3 only supports fieldless enums. PyO3 adds a class attribute for each variant, so you can access them in Python without defining `#[new]`. PyO3 also provides default implementations of `__richcmp__` and `__int__`, so they can be compared using `==`:
+
+```rust
+# use pyo3::prelude::*;
+#[pyclass]
+enum MyEnum {
+    Variant,
+    OtherVariant,
+}
+
+Python::with_gil(|py| {
+    let x = Py::new(py, MyEnum::Variant).unwrap();
+    let y = Py::new(py, MyEnum::OtherVariant).unwrap();
+    let cls = py.get_type::<MyEnum>();
+    pyo3::py_run!(py, x y cls, r#"
+        assert x == cls.Variant
+        assert y == cls.OtherVariant
+        assert x != y
+    "#)
+})
+```
+
+You can also convert your enums into `int`:
+
+```rust
+# use pyo3::prelude::*;
+#[pyclass]
+enum MyEnum {
+    Variant,
+    OtherVariant = 10,
+}
+
+Python::with_gil(|py| {
+    let cls = py.get_type::<MyEnum>();
+    let x = MyEnum::Variant as i32; // The exact value is assigned by the compiler.
+    pyo3::py_run!(py, cls x, r#"
+        assert int(cls.Variant) == x
+        assert int(cls.OtherVariant) == 10
+        assert cls.OtherVariant == 10  # You can also compare against int.
+        assert 10 == cls.OtherVariant
+    "#)
+})
+```
+
+PyO3 also provides `__repr__` for enums:
+
+```rust
+# use pyo3::prelude::*;
+#[pyclass]
+enum MyEnum{
+    Variant,
+    OtherVariant,
+}
+
+Python::with_gil(|py| {
+    let cls = py.get_type::<MyEnum>();
+    let x = Py::new(py, MyEnum::Variant).unwrap();
+    pyo3::py_run!(py, cls x, r#"
+        assert repr(x) == 'MyEnum.Variant'
+        assert repr(cls.OtherVariant) == 'MyEnum.OtherVariant'
+    "#)
+})
+```
+
+All methods defined by PyO3 can be overriden. For example here's how you override `__repr__`:
+
+```rust
+# use pyo3::prelude::*;
+#[pyclass]
+enum MyEnum {
+    Answer = 42,
+}
+
+#[pymethods]
+impl MyEnum {
+    fn __repr__(&self) -> &'static str {
+        "42"
+    }
+}
+
+Python::with_gil(|py| {
+    let cls = py.get_type::<MyEnum>();
+    pyo3::py_run!(py, cls, "assert repr(cls.Answer) == '42'")
+})
+```
+
+You may not use enums as a base class or let enums inherit from other classes.
+
+```rust,compile_fail
+# use pyo3::prelude::*;
+#[pyclass(subclass)]
+enum BadBase{
+    Var1,
+}
+```
+
+```rust,compile_fail
+# use pyo3::prelude::*;
+
+#[pyclass(subclass)]
+struct Base;
+
+#[pyclass(extends=Base)]
+enum BadSubclass{
+    Var1,
+}
+```
+
+`#[pyclass]` enums are currently not interoperable with `IntEnum` in Python.
+
 ## Implementation details
 
 The `#[pyclass]` macros rely on a lot of conditional code generation: each `#[pyclass]` can optionally have a `#[pymethods]` block as well as several different possible `#[pyproto]` trait implementations.

pyo3-macros-backend/src/pyclass.rs

@@ -407,6 +407,54 @@ struct PyClassEnumVariant<'a> {
     /* currently have no more options */
 }
 
+struct PyClassEnum<'a> {
+    ident: &'a syn::Ident,
+    // The underlying #[repr] of the enum, used to implement __int__ and __richcmp__.
+    // This matters when the underlying representation may not fit in `isize`.
+    repr_type: syn::Ident,
+    variants: Vec<PyClassEnumVariant<'a>>,
+}
+
+impl<'a> PyClassEnum<'a> {
+    fn new(enum_: &'a syn::ItemEnum) -> syn::Result<Self> {
+        fn is_numeric_type(t: &syn::Ident) -> bool {
+            [
+                "u8", "i8", "u16", "i16", "u32", "i32", "u64", "i64", "u128", "i128", "usize",
+                "isize",
+            ]
+            .iter()
+            .any(|&s| t == s)
+        }
+        let ident = &enum_.ident;
+        // According to the [reference](https://doc.rust-lang.org/reference/items/enumerations.html),
+        // "Under the default representation, the specified discriminant is interpreted as an isize
+        // value", so `isize` should be enough by default.
+        let mut repr_type = syn::Ident::new("isize", proc_macro2::Span::call_site());
+        if let Some(attr) = enum_.attrs.iter().find(|attr| attr.path.is_ident("repr")) {
+            let args =
+                attr.parse_args_with(Punctuated::<TokenStream, Token![!]>::parse_terminated)?;
+            if let Some(ident) = args
+                .into_iter()
+                .filter_map(|ts| syn::parse2::<syn::Ident>(ts).ok())
+                .find(is_numeric_type)
+            {
+                repr_type = ident;
+            }
+        }
+
+        let variants = enum_
+            .variants
+            .iter()
+            .map(extract_variant_data)
+            .collect::<syn::Result<_>>()?;
+        Ok(Self {
+            ident,
+            repr_type,
+            variants,
+        })
+    }
+}
+
 pub fn build_py_enum(
     enum_: &mut syn::ItemEnum,
     args: &PyClassArgs,
@@ -417,41 +465,37 @@ pub fn build_py_enum(
     if enum_.variants.is_empty() {
         bail_spanned!(enum_.brace_token.span => "Empty enums can't be #[pyclass].");
     }
-    let variants: Vec<PyClassEnumVariant> = enum_
-        .variants
-        .iter()
-        .map(extract_variant_data)
-        .collect::<syn::Result<_>>()?;
-    impl_enum(enum_, args, variants, method_type, options)
-}
-
-fn impl_enum(
-    enum_: &syn::ItemEnum,
-    args: &PyClassArgs,
-    variants: Vec<PyClassEnumVariant>,
-    methods_type: PyClassMethodsType,
-    options: PyClassPyO3Options,
-) -> syn::Result<TokenStream> {
-    let enum_name = &enum_.ident;
     let doc = utils::get_doc(
         &enum_.attrs,
         options
             .text_signature
             .as_ref()
             .map(|attr| (get_class_python_name(&enum_.ident, args), attr)),
     );
+    let enum_ = PyClassEnum::new(enum_)?;
+    impl_enum(enum_, args, doc, method_type, options)
+}
+
+fn impl_enum(
+    enum_: PyClassEnum,
+    args: &PyClassArgs,
+    doc: PythonDoc,
+    methods_type: PyClassMethodsType,
+    options: PyClassPyO3Options,
+) -> syn::Result<TokenStream> {
     let krate = get_pyo3_crate(&options.krate);
-    impl_enum_class(enum_name, args, variants, doc, methods_type, krate)
+    impl_enum_class(enum_, args, doc, methods_type, krate)
 }
 
 fn impl_enum_class(
-    cls: &syn::Ident,
+    enum_: PyClassEnum,
     args: &PyClassArgs,
-    variants: Vec<PyClassEnumVariant>,
     doc: PythonDoc,
     methods_type: PyClassMethodsType,
     krate: syn::Path,
 ) -> syn::Result<TokenStream> {
+    let cls = enum_.ident;
+    let variants = enum_.variants;
     let pytypeinfo = impl_pytypeinfo(cls, args, None);
     let pyclass_impls = PyClassImplsBuilder::new(
         cls,
@@ -476,13 +520,73 @@ fn impl_enum_class(
             fn __pyo3__repr__(&self) -> &'static str {
                 match self {
                     #(#variants_repr)*
-                    _ => unreachable!("Unsupported variant type."),
                 }
             }
         }
     };
 
-    let default_impls = gen_default_items(cls, vec![default_repr_impl]);
+    let repr_type = &enum_.repr_type;
+
+    let default_int = {
+        // This implementation allows us to convert &T to #repr_type without implementing `Copy`
+        let variants_to_int = variants.iter().map(|variant| {
+            let variant_name = variant.ident;
+            quote! { #cls::#variant_name => #cls::#variant_name as #repr_type, }
+        });
+        quote! {
+            #[doc(hidden)]
+            #[allow(non_snake_case)]
+            #[pyo3(name = "__int__")]
+            fn __pyo3__int__(&self) -> #repr_type {
+                match self {
+                    #(#variants_to_int)*
+                }
+            }
+        }
+    };
+
+    let default_richcmp = {
+        let variants_eq = variants.iter().map(|variant| {
+            let variant_name = variant.ident;
+            quote! {
+                (#cls::#variant_name, #cls::#variant_name) =>
+                    Ok(true.to_object(py)),
+            }
+        });
+        quote! {
+            #[doc(hidden)]
+            #[allow(non_snake_case)]
+            #[pyo3(name = "__richcmp__")]
+            fn __pyo3__richcmp__(
+                &self,
+                py: _pyo3::Python,
+                other: &_pyo3::PyAny,
+                op: _pyo3::basic::CompareOp
+            ) -> _pyo3::PyResult<_pyo3::PyObject> {
+                use _pyo3::conversion::ToPyObject;
+                use ::core::result::Result::*;
+                match op {
+                    _pyo3::basic::CompareOp::Eq => {
+                        if let Ok(i) = other.extract::<#repr_type>() {
+                            let self_val = self.__pyo3__int__();
+                            return Ok((self_val == i).to_object(py));
+                        }
+                        let other = other.extract::<_pyo3::PyRef<Self>>()?;
+                        let other = &*other;
+                        match (self, other) {
+                            #(#variants_eq)*
+                            _ => Ok(false.to_object(py)),
+                        }
+                    }
+                    _ => Ok(py.NotImplemented()),
+                }
+            }
+        }
+    };
+
+    let default_items =
+        gen_default_items(cls, vec![default_repr_impl, default_richcmp, default_int]);
+
     Ok(quote! {
         const _: () = {
             use #krate as _pyo3;
@@ -491,7 +595,7 @@ fn impl_enum_class(
 
             #pyclass_impls
 
-            #default_impls
+            #default_items
         };
     })
 }
@@ -527,9 +631,6 @@ fn extract_variant_data(variant: &syn::Variant) -> syn::Result<PyClassEnumVarian
         Fields::Unit => &variant.ident,
         _ => bail_spanned!(variant.span() => "Currently only support unit variants."),
     };
-    if let Some(discriminant) = variant.discriminant.as_ref() {
-        bail_spanned!(discriminant.0.span() => "Currently does not support discriminats.")
-    };
     Ok(PyClassEnumVariant { ident })
 }