1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
//! This crate declares only the proc macro attributes, as a crate defining proc macro attributes
//! must not contain any other public items.

#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use pyo3_macros_backend::{
    build_derive_from_pyobject, build_py_class, build_py_enum, build_py_function, build_py_methods,
    pymodule_function_impl, pymodule_module_impl, PyClassArgs, PyClassMethodsType,
    PyFunctionOptions,
};
use quote::quote;
use syn::{parse::Nothing, parse_macro_input, Item};

/// A proc macro used to implement Python modules.
///
/// The name of the module will be taken from the function name, unless `#[pyo3(name = "my_name")]`
/// is also annotated on the function to override the name. **Important**: the module name should
/// match the `lib.name` setting in `Cargo.toml`, so that Python is able to import the module
/// without needing a custom import loader.
///
/// Functions annotated with `#[pymodule]` can also be annotated with the following:
///
/// |  Annotation  |  Description |
/// | :-  | :- |
/// | `#[pyo3(name = "...")]` | Defines the name of the module in Python. |
///
/// For more on creating Python modules see the [module section of the guide][1].
///
/// Due to technical limitations on how `#[pymodule]` is implemented, a function marked
/// `#[pymodule]` cannot have a module with the same name in the same scope. (The
/// `#[pymodule]` implementation generates a hidden module with the same name containing
/// metadata about the module, which is used by `wrap_pymodule!`).
///
/// [1]: https://pyo3.rs/latest/module.html
#[proc_macro_attribute]
pub fn pymodule(args: TokenStream, input: TokenStream) -> TokenStream {
    parse_macro_input!(args as Nothing);
    match parse_macro_input!(input as Item) {
        Item::Mod(module) => if cfg!(feature = "experimental-declarative-modules") {
            pymodule_module_impl(module)
        } else {
            Err(syn::Error::new_spanned(
                module,
                "#[pymodule] requires the 'experimental-declarative-modules' feature to be used on Rust modules.",
            ))
        },
        Item::Fn(function) => pymodule_function_impl(function),
        unsupported => Err(syn::Error::new_spanned(
            unsupported,
            "#[pymodule] only supports modules and functions.",
        )),
    }
    .unwrap_or_compile_error()
    .into()
}

#[proc_macro_attribute]
pub fn pyclass(attr: TokenStream, input: TokenStream) -> TokenStream {
    let item = parse_macro_input!(input as Item);
    match item {
        Item::Struct(struct_) => pyclass_impl(attr, struct_, methods_type()),
        Item::Enum(enum_) => pyclass_enum_impl(attr, enum_, methods_type()),
        unsupported => {
            syn::Error::new_spanned(unsupported, "#[pyclass] only supports structs and enums.")
                .into_compile_error()
                .into()
        }
    }
}

/// A proc macro used to expose methods to Python.
///
/// Methods within a `#[pymethods]` block can be annotated with  as well as the following:
///
/// |  Annotation  |  Description |
/// | :-  | :- |
/// | [`#[new]`][4]  | Defines the class constructor, like Python's `__new__` method. |
/// | [`#[getter]`][5] and [`#[setter]`][5] | These define getters and setters, similar to Python's `@property` decorator. This is useful for getters/setters that require computation or side effects; if that is not the case consider using [`#[pyo3(get, set)]`][11] on the struct's field(s).|
/// | [`#[staticmethod]`][6]| Defines the method as a staticmethod, like Python's `@staticmethod` decorator.|
/// | [`#[classmethod]`][7]  | Defines the method as a classmethod, like Python's `@classmethod` decorator.|
/// | [`#[classattr]`][9]  | Defines a class variable. |
/// | [`#[args]`][10]  | Deprecated way to define a method's default arguments and allows the function to receive `*args` and `**kwargs`. Use `#[pyo3(signature = (...))]` instead. |
/// | <nobr>[`#[pyo3(<option> = <value>)`][pyo3-method-options]</nobr> | Any of the `#[pyo3]` options supported on [`macro@pyfunction`]. |
///
/// For more on creating class methods,
/// see the [class section of the guide][1].
///
/// If the [`multiple-pymethods`][2] feature is enabled, it is possible to implement
/// multiple `#[pymethods]` blocks for a single `#[pyclass]`.
/// This will add a transitive dependency on the [`inventory`][3] crate.
///
/// [1]: https://pyo3.rs/latest/class.html#instance-methods
/// [2]: https://pyo3.rs/latest/features.html#multiple-pymethods
/// [3]: https://docs.rs/inventory/
/// [4]: https://pyo3.rs/latest/class.html#constructor
/// [5]: https://pyo3.rs/latest/class.html#object-properties-using-getter-and-setter
/// [6]: https://pyo3.rs/latest/class.html#static-methods
/// [7]: https://pyo3.rs/latest/class.html#class-methods
/// [8]: https://pyo3.rs/latest/class.html#callable-objects
/// [9]: https://pyo3.rs/latest/class.html#class-attributes
/// [10]: https://pyo3.rs/latest/class.html#method-arguments
/// [11]: https://pyo3.rs/latest/class.html#object-properties-using-pyo3get-set
#[proc_macro_attribute]
pub fn pymethods(attr: TokenStream, input: TokenStream) -> TokenStream {
    let methods_type = if cfg!(feature = "multiple-pymethods") {
        PyClassMethodsType::Inventory
    } else {
        PyClassMethodsType::Specialization
    };
    pymethods_impl(attr, input, methods_type)
}

/// A proc macro used to expose Rust functions to Python.
///
/// Functions annotated with `#[pyfunction]` can also be annotated with the following `#[pyo3]`
/// options:
///
/// |  Annotation  |  Description |
/// | :-  | :- |
/// | `#[pyo3(name = "...")]` | Defines the name of the function in Python. |
/// | `#[pyo3(text_signature = "...")]` | Defines the `__text_signature__` attribute of the function in Python. |
/// | `#[pyo3(pass_module)]` | Passes the module containing the function as a `&PyModule` first argument to the function. |
///
/// For more on exposing functions see the [function section of the guide][1].
///
/// Due to technical limitations on how `#[pyfunction]` is implemented, a function marked
/// `#[pyfunction]` cannot have a module with the same name in the same scope. (The
/// `#[pyfunction]` implementation generates a hidden module with the same name containing
/// metadata about the function, which is used by `wrap_pyfunction!`).
///
/// [1]: https://pyo3.rs/latest/function.html
#[proc_macro_attribute]
pub fn pyfunction(attr: TokenStream, input: TokenStream) -> TokenStream {
    let mut ast = parse_macro_input!(input as syn::ItemFn);
    let options = parse_macro_input!(attr as PyFunctionOptions);

    let expanded = build_py_function(&mut ast, options).unwrap_or_compile_error();

    quote!(
        #ast
        #expanded
    )
    .into()
}

#[proc_macro_derive(FromPyObject, attributes(pyo3))]
pub fn derive_from_py_object(item: TokenStream) -> TokenStream {
    let ast = parse_macro_input!(item as syn::DeriveInput);
    let expanded = build_derive_from_pyobject(&ast).unwrap_or_compile_error();
    quote!(
        #expanded
    )
    .into()
}

fn pyclass_impl(
    attrs: TokenStream,
    mut ast: syn::ItemStruct,
    methods_type: PyClassMethodsType,
) -> TokenStream {
    let args = parse_macro_input!(attrs with PyClassArgs::parse_stuct_args);
    let expanded = build_py_class(&mut ast, args, methods_type).unwrap_or_compile_error();

    quote!(
        #ast
        #expanded
    )
    .into()
}

fn pyclass_enum_impl(
    attrs: TokenStream,
    mut ast: syn::ItemEnum,
    methods_type: PyClassMethodsType,
) -> TokenStream {
    let args = parse_macro_input!(attrs with PyClassArgs::parse_enum_args);
    let expanded = build_py_enum(&mut ast, args, methods_type).unwrap_or_compile_error();

    quote!(
        #ast
        #expanded
    )
    .into()
}

fn pymethods_impl(
    attr: TokenStream,
    input: TokenStream,
    methods_type: PyClassMethodsType,
) -> TokenStream {
    let mut ast = parse_macro_input!(input as syn::ItemImpl);
    // Apply all options as a #[pyo3] attribute on the ItemImpl
    // e.g. #[pymethods(crate = "crate")] impl Foo { }
    // -> #[pyo3(crate = "crate")] impl Foo { }
    let attr: TokenStream2 = attr.into();
    ast.attrs.push(syn::parse_quote!( #[pyo3(#attr)] ));
    let expanded = build_py_methods(&mut ast, methods_type).unwrap_or_compile_error();

    quote!(
        #ast
        #expanded
    )
    .into()
}

fn methods_type() -> PyClassMethodsType {
    if cfg!(feature = "multiple-pymethods") {
        PyClassMethodsType::Inventory
    } else {
        PyClassMethodsType::Specialization
    }
}

trait UnwrapOrCompileError {
    fn unwrap_or_compile_error(self) -> TokenStream2;
}

impl UnwrapOrCompileError for syn::Result<TokenStream2> {
    fn unwrap_or_compile_error(self) -> TokenStream2 {
        self.unwrap_or_else(|e| e.into_compile_error())
    }
}