module.rs

// Copyright (c) 2017-present PyO3 Project and Contributors
//! Code generation for the function that initializes a python module and adds classes and function.

use crate::{
    attributes::{
        self, is_attribute_ident, take_attributes, take_pyo3_options, NameAttribute,
        PyO3PathAttribute,
    },
    pyfunction::{impl_wrap_pyfunction, PyFunctionOptions},
    utils::{get_pyo3_path, PythonDoc},
};
use proc_macro2::{Span, TokenStream};
use quote::quote;
use syn::{
    ext::IdentExt,
    parse::{Parse, ParseStream},
    spanned::Spanned,
    token::Comma,
    Ident, Path, Result,
};

#[derive(Default)]
pub struct PyModuleOptions {
    pyo3_path: Option<PyO3PathAttribute>,
    name: Option<syn::Ident>,
}

impl PyModuleOptions {
    pub fn from_attrs(attrs: &mut Vec<syn::Attribute>) -> Result<Self> {
        let mut options: PyModuleOptions = Default::default();

        for option in take_pyo3_options(attrs)? {
            match option {
                PyModulePyO3Option::Name(name) => options.set_name(name.0)?,
                PyModulePyO3Option::PyO3Path(path) => options.set_pyo3_path(path)?,
            }
        }

        Ok(options)
    }

    fn set_name(&mut self, name: syn::Ident) -> Result<()> {
        ensure_spanned!(
            self.name.is_none(),
            name.span() => "`name` may only be specified once"
        );

        self.name = Some(name);
        Ok(())
    }

    fn set_pyo3_path(&mut self, path: PyO3PathAttribute) -> Result<()> {
        ensure_spanned!(
            self.pyo3_path.is_none(),
            path.0.span() => "`pyo3_path` may only be specified once"
        );

        self.pyo3_path = Some(path);
        Ok(())
    }
}

/// Generates the function that is called by the python interpreter to initialize the native
/// module
pub fn py_init(fnname: &Ident, options: PyModuleOptions, doc: PythonDoc) -> TokenStream {
    let name = options.name.unwrap_or_else(|| fnname.unraw());
    let pyo3_path = get_pyo3_path(&options.pyo3_path);
    let cb_name = Ident::new(&format!("PyInit_{}", name), Span::call_site());

    quote! {
        #[no_mangle]
        #[allow(non_snake_case)]
        /// This autogenerated function is called by the python interpreter when importing
        /// the module.
        pub unsafe extern "C" fn #cb_name() -> *mut #pyo3_path::ffi::PyObject {
            use #pyo3_path as _pyo3;
            use _pyo3::derive_utils::ModuleDef;
            static NAME: &str = concat!(stringify!(#name), "\0");
            static DOC: &str = #doc;
            static MODULE_DEF: ModuleDef = unsafe { ModuleDef::new(NAME, DOC) };

            _pyo3::callback::handle_panic(|_py| { MODULE_DEF.make_module(_py, #fnname) })
        }
    }
}

/// Finds and takes care of the #[pyfn(...)] in `#[pymodule]`
pub fn process_functions_in_module(func: &mut syn::ItemFn) -> syn::Result<()> {
    let mut stmts: Vec<syn::Stmt> = Vec::new();

    for mut stmt in func.block.stmts.drain(..) {
        if let syn::Stmt::Item(syn::Item::Fn(func)) = &mut stmt {
            if let Some(pyfn_args) = get_pyfn_attr(&mut func.attrs)? {
                let module_name = pyfn_args.modname;
                let (ident, wrapped_function) = impl_wrap_pyfunction(func, pyfn_args.options)?;
                let statements: Vec<syn::Stmt> = syn::parse_quote! {
                    #wrapped_function
                    #module_name.add_function(#ident(#module_name)?)?;
                };
                stmts.extend(statements);
            }
        };
        stmts.push(stmt);
    }

    func.block.stmts = stmts;
    Ok(())
}

pub struct PyFnArgs {
    modname: Path,
    options: PyFunctionOptions,
}

impl Parse for PyFnArgs {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let modname = input.parse().map_err(
            |e| err_spanned!(e.span() => "expected module as first argument to #[pyfn()]"),
        )?;

        if input.is_empty() {
            return Ok(Self {
                modname,
                options: Default::default(),
            });
        }

        let _: Comma = input.parse()?;

        Ok(Self {
            modname,
            options: input.parse()?,
        })
    }
}

/// Extracts the data from the #[pyfn(...)] attribute of a function
fn get_pyfn_attr(attrs: &mut Vec<syn::Attribute>) -> syn::Result<Option<PyFnArgs>> {
    let mut pyfn_args: Option<PyFnArgs> = None;

    take_attributes(attrs, |attr| {
        if is_attribute_ident(attr, "pyfn") {
            ensure_spanned!(
                pyfn_args.is_none(),
                attr.span() => "`#[pyfn] may only be specified once"
            );
            pyfn_args = Some(attr.parse_args()?);
            Ok(true)
        } else {
            Ok(false)
        }
    })?;

    if let Some(pyfn_args) = &mut pyfn_args {
        pyfn_args
            .options
            .add_attributes(take_pyo3_options(attrs)?)?;
    }

    Ok(pyfn_args)
}

enum PyModulePyO3Option {
    PyO3Path(PyO3PathAttribute),
    Name(NameAttribute),
}

impl Parse for PyModulePyO3Option {
    fn parse(input: ParseStream) -> Result<Self> {
        let lookahead = input.lookahead1();
        if lookahead.peek(attributes::kw::name) {
            input.parse().map(PyModulePyO3Option::Name)
        } else if lookahead.peek(attributes::kw::pyo3_path) {
            input.parse().map(PyModulePyO3Option::PyO3Path)
        } else {
            Err(lookahead.error())
        }
    }
}