pyo3_macros_backend/
pyfunction.rs

1use crate::attributes::KeywordAttribute;
2use crate::combine_errors::CombineErrors;
3#[cfg(feature = "experimental-inspect")]
4use crate::introspection::{function_introspection_code, introspection_id_const};
5#[cfg(feature = "experimental-inspect")]
6use crate::py_expr::PyExpr;
7#[cfg(feature = "experimental-inspect")]
8use crate::utils::get_doc;
9use crate::utils::Ctx;
10use crate::{
11    attributes::{
12        self, get_pyo3_options, take_attributes, take_pyo3_options, CrateAttribute,
13        FromPyWithAttribute, NameAttribute, TextSignatureAttribute,
14    },
15    method::{self, CallingConvention, ClassMethodReceiver, FnArg, SelfConversionPolicy},
16    pymethod::check_generic,
17};
18use proc_macro2::{Span, TokenStream};
19use quote::{format_ident, quote, ToTokens};
20use std::cmp::PartialEq;
21use std::ffi::CString;
22#[cfg(feature = "experimental-inspect")]
23use std::iter::empty;
24use syn::parse::{Parse, ParseStream};
25use syn::punctuated::Punctuated;
26#[cfg(feature = "experimental-inspect")]
27use syn::ReturnType;
28use syn::{ext::IdentExt, spanned::Spanned, LitCStr, LitStr, Path, Result, Token};
29
30mod signature;
31
32pub use self::signature::{ConstructorAttribute, FunctionSignature, SignatureAttribute};
33
34#[derive(Clone, Debug)]
35pub struct PyFunctionArgPyO3Attributes {
36    pub from_py_with: Option<FromPyWithAttribute>,
37    pub cancel_handle: Option<attributes::kw::cancel_handle>,
38}
39
40enum PyFunctionArgPyO3Attribute {
41    FromPyWith(FromPyWithAttribute),
42    CancelHandle(attributes::kw::cancel_handle),
43}
44
45impl Parse for PyFunctionArgPyO3Attribute {
46    fn parse(input: ParseStream<'_>) -> Result<Self> {
47        let lookahead = input.lookahead1();
48        if lookahead.peek(attributes::kw::cancel_handle) {
49            input.parse().map(PyFunctionArgPyO3Attribute::CancelHandle)
50        } else if lookahead.peek(attributes::kw::from_py_with) {
51            input.parse().map(PyFunctionArgPyO3Attribute::FromPyWith)
52        } else {
53            Err(lookahead.error())
54        }
55    }
56}
57
58impl PyFunctionArgPyO3Attributes {
59    /// Parses #[pyo3(from_python_with = "func")]
60    pub fn from_attrs(attrs: &mut Vec<syn::Attribute>) -> syn::Result<Self> {
61        let mut attributes = PyFunctionArgPyO3Attributes {
62            from_py_with: None,
63            cancel_handle: None,
64        };
65        take_attributes(attrs, |attr| {
66            if let Some(pyo3_attrs) = get_pyo3_options(attr)? {
67                for attr in pyo3_attrs {
68                    match attr {
69                        PyFunctionArgPyO3Attribute::FromPyWith(from_py_with) => {
70                            ensure_spanned!(
71                                attributes.from_py_with.is_none(),
72                                from_py_with.span() => "`from_py_with` may only be specified once per argument"
73                            );
74                            attributes.from_py_with = Some(from_py_with);
75                        }
76                        PyFunctionArgPyO3Attribute::CancelHandle(cancel_handle) => {
77                            ensure_spanned!(
78                                attributes.cancel_handle.is_none(),
79                                cancel_handle.span() => "`cancel_handle` may only be specified once per argument"
80                            );
81                            attributes.cancel_handle = Some(cancel_handle);
82                        }
83                    }
84                    ensure_spanned!(
85                        attributes.from_py_with.is_none() || attributes.cancel_handle.is_none(),
86                        attributes.cancel_handle.unwrap().span() => "`from_py_with` and `cancel_handle` cannot be specified together"
87                    );
88                }
89                Ok(true)
90            } else {
91                Ok(false)
92            }
93        })?;
94        Ok(attributes)
95    }
96}
97
98type PyFunctionWarningMessageAttribute = KeywordAttribute<attributes::kw::message, LitStr>;
99type PyFunctionWarningCategoryAttribute = KeywordAttribute<attributes::kw::category, Path>;
100
101pub struct PyFunctionWarningAttribute {
102    pub message: PyFunctionWarningMessageAttribute,
103    pub category: Option<PyFunctionWarningCategoryAttribute>,
104    pub span: Span,
105}
106
107#[derive(PartialEq, Clone)]
108pub enum PyFunctionWarningCategory {
109    Path(Path),
110    UserWarning,
111    DeprecationWarning, // TODO: unused for now, intended for pyo3(deprecated) special-case
112}
113
114#[derive(Clone)]
115pub struct PyFunctionWarning {
116    pub message: LitStr,
117    pub category: PyFunctionWarningCategory,
118    pub span: Span,
119}
120
121impl From<PyFunctionWarningAttribute> for PyFunctionWarning {
122    fn from(value: PyFunctionWarningAttribute) -> Self {
123        Self {
124            message: value.message.value,
125            category: value
126                .category
127                .map_or(PyFunctionWarningCategory::UserWarning, |cat| {
128                    PyFunctionWarningCategory::Path(cat.value)
129                }),
130            span: value.span,
131        }
132    }
133}
134
135pub trait WarningFactory {
136    fn build_py_warning(&self, ctx: &Ctx) -> TokenStream;
137    fn span(&self) -> Span;
138}
139
140impl WarningFactory for PyFunctionWarning {
141    fn build_py_warning(&self, ctx: &Ctx) -> TokenStream {
142        let message = &self.message.value();
143        let c_message = LitCStr::new(
144            &CString::new(message.clone()).unwrap(),
145            Spanned::span(&message),
146        );
147        let pyo3_path = &ctx.pyo3_path;
148        let category = match &self.category {
149            PyFunctionWarningCategory::Path(path) => quote! {#path},
150            PyFunctionWarningCategory::UserWarning => {
151                quote! {#pyo3_path::exceptions::PyUserWarning}
152            }
153            PyFunctionWarningCategory::DeprecationWarning => {
154                quote! {#pyo3_path::exceptions::PyDeprecationWarning}
155            }
156        };
157        quote! {
158            #pyo3_path::PyErr::warn(py, &<#category as #pyo3_path::PyTypeInfo>::type_object(py), #c_message, 1)?;
159        }
160    }
161
162    fn span(&self) -> Span {
163        self.span
164    }
165}
166
167impl<T: WarningFactory> WarningFactory for Vec<T> {
168    fn build_py_warning(&self, ctx: &Ctx) -> TokenStream {
169        let warnings = self.iter().map(|warning| warning.build_py_warning(ctx));
170
171        quote! {
172            #(#warnings)*
173        }
174    }
175
176    fn span(&self) -> Span {
177        self.iter()
178            .map(|val| val.span())
179            .reduce(|acc, span| acc.join(span).unwrap_or(acc))
180            .unwrap()
181    }
182}
183
184impl Parse for PyFunctionWarningAttribute {
185    fn parse(input: ParseStream<'_>) -> Result<Self> {
186        let mut message: Option<PyFunctionWarningMessageAttribute> = None;
187        let mut category: Option<PyFunctionWarningCategoryAttribute> = None;
188
189        let span = input.parse::<attributes::kw::warn>()?.span();
190
191        let content;
192        syn::parenthesized!(content in input);
193
194        while !content.is_empty() {
195            let lookahead = content.lookahead1();
196
197            if lookahead.peek(attributes::kw::message) {
198                message = content
199                    .parse::<PyFunctionWarningMessageAttribute>()
200                    .map(Some)?;
201            } else if lookahead.peek(attributes::kw::category) {
202                category = content
203                    .parse::<PyFunctionWarningCategoryAttribute>()
204                    .map(Some)?;
205            } else {
206                return Err(lookahead.error());
207            }
208
209            if content.peek(Token![,]) {
210                content.parse::<Token![,]>()?;
211            }
212        }
213
214        Ok(PyFunctionWarningAttribute {
215            message: message.ok_or(syn::Error::new(
216                content.span(),
217                "missing `message` in `warn` attribute",
218            ))?,
219            category,
220            span,
221        })
222    }
223}
224
225impl ToTokens for PyFunctionWarningAttribute {
226    fn to_tokens(&self, tokens: &mut TokenStream) {
227        let message_tokens = self.message.to_token_stream();
228        let category_tokens = self
229            .category
230            .as_ref()
231            .map_or(quote! {}, |cat| cat.to_token_stream());
232
233        let token_stream = quote! {
234            warn(#message_tokens, #category_tokens)
235        };
236
237        tokens.extend(token_stream);
238    }
239}
240
241#[derive(Default)]
242pub struct PyFunctionOptions {
243    pub pass_module: Option<attributes::kw::pass_module>,
244    pub name: Option<NameAttribute>,
245    pub signature: Option<SignatureAttribute>,
246    pub text_signature: Option<TextSignatureAttribute>,
247    pub krate: Option<CrateAttribute>,
248    pub warnings: Vec<PyFunctionWarning>,
249}
250
251impl Parse for PyFunctionOptions {
252    fn parse(input: ParseStream<'_>) -> Result<Self> {
253        let mut options = PyFunctionOptions::default();
254
255        let attrs = Punctuated::<PyFunctionOption, syn::Token![,]>::parse_terminated(input)?;
256        options.add_attributes(attrs)?;
257
258        Ok(options)
259    }
260}
261
262pub enum PyFunctionOption {
263    Name(NameAttribute),
264    PassModule(attributes::kw::pass_module),
265    Signature(SignatureAttribute),
266    TextSignature(TextSignatureAttribute),
267    Crate(CrateAttribute),
268    Warning(PyFunctionWarningAttribute),
269}
270
271impl Parse for PyFunctionOption {
272    fn parse(input: ParseStream<'_>) -> Result<Self> {
273        let lookahead = input.lookahead1();
274        if lookahead.peek(attributes::kw::name) {
275            input.parse().map(PyFunctionOption::Name)
276        } else if lookahead.peek(attributes::kw::pass_module) {
277            input.parse().map(PyFunctionOption::PassModule)
278        } else if lookahead.peek(attributes::kw::signature) {
279            input.parse().map(PyFunctionOption::Signature)
280        } else if lookahead.peek(attributes::kw::text_signature) {
281            input.parse().map(PyFunctionOption::TextSignature)
282        } else if lookahead.peek(syn::Token![crate]) {
283            input.parse().map(PyFunctionOption::Crate)
284        } else if lookahead.peek(attributes::kw::warn) {
285            input.parse().map(PyFunctionOption::Warning)
286        } else {
287            Err(lookahead.error())
288        }
289    }
290}
291
292impl PyFunctionOptions {
293    pub fn from_attrs(attrs: &mut Vec<syn::Attribute>) -> syn::Result<Self> {
294        let mut options = PyFunctionOptions::default();
295        options.add_attributes(take_pyo3_options(attrs)?)?;
296        Ok(options)
297    }
298
299    pub fn add_attributes(
300        &mut self,
301        attrs: impl IntoIterator<Item = PyFunctionOption>,
302    ) -> Result<()> {
303        macro_rules! set_option {
304            ($key:ident) => {
305                {
306                    ensure_spanned!(
307                        self.$key.is_none(),
308                        $key.span() => concat!("`", stringify!($key), "` may only be specified once")
309                    );
310                    self.$key = Some($key);
311                }
312            };
313        }
314        for attr in attrs {
315            match attr {
316                PyFunctionOption::Name(name) => set_option!(name),
317                PyFunctionOption::PassModule(pass_module) => set_option!(pass_module),
318                PyFunctionOption::Signature(signature) => set_option!(signature),
319                PyFunctionOption::TextSignature(text_signature) => set_option!(text_signature),
320                PyFunctionOption::Crate(krate) => set_option!(krate),
321                PyFunctionOption::Warning(warning) => {
322                    self.warnings.push(warning.into());
323                }
324            }
325        }
326        Ok(())
327    }
328}
329
330pub fn build_py_function(
331    ast: &mut syn::ItemFn,
332    mut options: PyFunctionOptions,
333) -> syn::Result<TokenStream> {
334    options.add_attributes(take_pyo3_options(&mut ast.attrs)?)?;
335    impl_wrap_pyfunction(ast, options)
336}
337
338/// Generates python wrapper over a function that allows adding it to a python module as a python
339/// function
340pub fn impl_wrap_pyfunction(
341    func: &mut syn::ItemFn,
342    options: PyFunctionOptions,
343) -> syn::Result<TokenStream> {
344    check_generic(&func.sig)?;
345    let PyFunctionOptions {
346        pass_module,
347        name,
348        signature,
349        text_signature,
350        krate,
351        warnings,
352    } = options;
353
354    let ctx = &Ctx::new(&krate, Some(&func.sig));
355    let Ctx { pyo3_path, .. } = &ctx;
356
357    let python_name = name
358        .as_ref()
359        .map_or_else(|| &func.sig.ident, |name| &name.value.0)
360        .unraw();
361
362    let tp = if pass_module.is_some() {
363        let span = match func.sig.inputs.first() {
364            Some(syn::FnArg::Typed(first_arg)) => first_arg.ty.span(),
365            Some(syn::FnArg::Receiver(_)) | None => bail_spanned!(
366                func.sig.paren_token.span.join() => "expected `&PyModule` or `Py<PyModule>` as first argument with `pass_module`"
367            ),
368        };
369        method::FnType::FnModule(span)
370    } else {
371        method::FnType::FnStatic
372    };
373
374    let arguments = func
375        .sig
376        .inputs
377        .iter_mut()
378        .skip(if tp.skip_first_rust_argument_in_python_signature() {
379            1
380        } else {
381            0
382        })
383        .map(FnArg::parse)
384        .try_combine_syn_errors()?;
385
386    let signature = if let Some(signature) = signature {
387        FunctionSignature::from_arguments_and_attribute(arguments, signature)?
388    } else {
389        FunctionSignature::from_arguments(arguments)
390    };
391
392    let spec = method::FnSpec {
393        tp,
394        name: &func.sig.ident,
395        python_name,
396        signature,
397        text_signature,
398        asyncness: func.sig.asyncness,
399        unsafety: func.sig.unsafety,
400        warnings,
401        output: func.sig.output.clone(),
402    };
403
404    let vis = &func.vis;
405    let name = &func.sig.ident;
406
407    #[cfg(feature = "experimental-inspect")]
408    let introspection = function_introspection_code(
409        pyo3_path,
410        Some(name),
411        // `name` is the Rust identifier, which `#[pyo3(name = "...")]` overrides for Python
412        &spec.python_name.to_string(),
413        &spec.signature,
414        None,
415        match &func.sig.output {
416            ReturnType::Type(_, t) => PyExpr::from_return_type((**t).clone(), None),
417            ReturnType::Default => PyExpr::none(),
418        },
419        empty(),
420        func.sig.asyncness.is_some(),
421        false,
422        get_doc(&func.attrs, None).as_ref(),
423        None,
424    );
425    #[cfg(not(feature = "experimental-inspect"))]
426    let introspection = quote! {};
427    #[cfg(feature = "experimental-inspect")]
428    let introspection_id = introspection_id_const();
429    #[cfg(not(feature = "experimental-inspect"))]
430    let introspection_id = quote! {};
431
432    let wrapper_ident = format_ident!("__pyfunction_{}", spec.name);
433    if spec.asyncness.is_some() {
434        ensure_spanned!(
435            cfg!(feature = "experimental-async"),
436            spec.asyncness.span() => "async functions are only supported with the `experimental-async` feature"
437        );
438    }
439    let calling_convention = CallingConvention::from_signature(&spec.signature);
440    let wrapper = spec.get_wrapper_function(
441        &wrapper_ident,
442        None,
443        calling_convention,
444        SelfConversionPolicy::checked(),
445        ClassMethodReceiver::Class,
446        ctx,
447    )?;
448    let methoddef = spec.get_methoddef(
449        wrapper_ident,
450        spec.get_doc(&func.attrs).as_ref(),
451        calling_convention,
452        ctx,
453    )?;
454
455    let wrapped_pyfunction = quote! {
456        // Create a module with the same name as the `#[pyfunction]` - this way `use <the function>`
457        // will actually bring both the module and the function into scope.
458        #[doc(hidden)]
459        #vis mod #name {
460            pub(crate) struct MakeDef;
461            pub static _PYO3_DEF: #pyo3_path::impl_::pyfunction::PyFunctionDef = MakeDef::_PYO3_DEF;
462            #introspection_id
463        }
464
465        // Generate the definition in the same scope as the original function -
466        // this avoids complications around the fact that the generated module has a different scope
467        // (and `super` doesn't always refer to the outer scope, e.g. if the `#[pyfunction] is
468        // inside a function body)
469        #[allow(unknown_lints, non_local_definitions)]
470        impl #name::MakeDef {
471            // We're using this to initialize a static, so it's fine.
472            #[allow(clippy::declare_interior_mutable_const)]
473            const _PYO3_DEF: #pyo3_path::impl_::pyfunction::PyFunctionDef =
474                #pyo3_path::impl_::pyfunction::PyFunctionDef::from_method_def(#methoddef);
475        }
476
477        #[allow(non_snake_case)]
478        #wrapper
479
480        #introspection
481    };
482    Ok(wrapped_pyfunction)
483}
pyo3_macros_backend/pyfunction.rs

pyo3_macros_backend/
pyfunction.rs
