pyo3_stub_gen_derive/
lib.rs

1mod gen_stub;
2
3use proc_macro::TokenStream;
4
5/// Embed metadata for Python stub file generation for `#[pyclass]` macro
6///
7/// ```
8/// #[pyo3_stub_gen_derive::gen_stub_pyclass]
9/// #[pyo3::pyclass(mapping, module = "my_module", name = "Placeholder")]
10/// #[derive(Debug, Clone)]
11/// pub struct PyPlaceholder {
12///     #[pyo3(get)]
13///     pub name: String,
14///     #[pyo3(get)]
15///     pub ndim: usize,
16///     #[pyo3(get)]
17///     pub description: Option<String>,
18///     pub custom_latex: Option<String>,
19/// }
20/// ```
21#[proc_macro_attribute]
22pub fn gen_stub_pyclass(_attr: TokenStream, item: TokenStream) -> TokenStream {
23    gen_stub::pyclass(item.into())
24        .unwrap_or_else(|err| err.to_compile_error())
25        .into()
26}
27
28/// Embed metadata for Python stub file generation for `#[pyclass]` macro with enum
29///
30/// ```
31/// #[pyo3_stub_gen_derive::gen_stub_pyclass_enum]
32/// #[pyo3::pyclass(module = "my_module", name = "DataType")]
33/// #[derive(Debug, Clone, PartialEq, Eq, Hash)]
34/// pub enum PyDataType {
35///     #[pyo3(name = "FLOAT")]
36///     Float,
37///     #[pyo3(name = "INTEGER")]
38///     Integer,
39/// }
40/// ```
41#[proc_macro_attribute]
42pub fn gen_stub_pyclass_enum(_attr: TokenStream, item: TokenStream) -> TokenStream {
43    gen_stub::pyclass_enum(item.into())
44        .unwrap_or_else(|err| err.to_compile_error())
45        .into()
46}
47
48/// Embed metadata for Python stub file generation for `#[pyclass]` macro with a complex enum
49///
50/// ```
51/// #[pyo3_stub_gen_derive::gen_stub_pyclass_complex_enum]
52/// #[pyo3::pyclass(module = "my_module", name = "DataType")]
53/// #[derive(Debug, Clone)]
54/// pub enum PyDataType {
55///     #[pyo3(name = "FLOAT")]
56///     Float{f: f64},
57///     #[pyo3(name = "INTEGER")]
58///     Integer(i64),
59/// }
60/// ```
61#[proc_macro_attribute]
62pub fn gen_stub_pyclass_complex_enum(_attr: TokenStream, item: TokenStream) -> TokenStream {
63    gen_stub::pyclass_complex_enum(item.into())
64        .unwrap_or_else(|err| err.to_compile_error())
65        .into()
66}
67
68/// Embed metadata for Python stub file generation for `#[pymethods]` macro
69///
70/// ```
71/// #[pyo3_stub_gen_derive::gen_stub_pyclass]
72/// #[pyo3::pyclass]
73/// struct A {}
74///
75/// #[pyo3_stub_gen_derive::gen_stub_pymethods]
76/// #[pyo3::pymethods]
77/// impl A {
78///     #[getter]
79///     fn f(&self) -> Vec<u32> {
80///        todo!()
81///     }
82/// }
83/// ```
84#[proc_macro_attribute]
85pub fn gen_stub_pymethods(_attr: TokenStream, item: TokenStream) -> TokenStream {
86    gen_stub::pymethods(item.into())
87        .unwrap_or_else(|err| err.to_compile_error())
88        .into()
89}
90
91/// Embed metadata for Python stub file generation for `#[pyfunction]` macro
92///
93/// ```
94/// #[pyo3_stub_gen_derive::gen_stub_pyfunction]
95/// #[pyo3::pyfunction]
96/// #[pyo3(name = "is_odd")]
97/// pub fn is_odd(x: u32) -> bool {
98///     todo!()
99/// }
100/// ```
101///
102/// The function attributed by `#[gen_stub_pyfunction]` will be appended to default stub file.
103/// If you want to append this function to another module, add `module` attribute.
104///
105/// ```
106/// #[pyo3_stub_gen_derive::gen_stub_pyfunction(module = "my_module.experimental")]
107/// #[pyo3::pyfunction]
108/// #[pyo3(name = "is_odd")]
109/// pub fn is_odd(x: u32) -> bool {
110///     todo!()
111/// }
112/// ```
113#[proc_macro_attribute]
114pub fn gen_stub_pyfunction(attr: TokenStream, item: TokenStream) -> TokenStream {
115    gen_stub::pyfunction(attr.into(), item.into())
116        .unwrap_or_else(|err| err.to_compile_error())
117        .into()
118}
119
120/// Do nothing but remove all `#[gen_stub(xxx)]` for `pyclass`, `pymethods`, and `pyfunction`.
121///
122/// It is useful to use `#[gen_stub(xxx)]` under feature-gating stub-gen.
123///
124/// E.g., only generate .pyi when `stub-gen` feature is turned-on:
125/// ```ignore
126/// #[cfg_attr(feature = "stub-gen", pyo3_stub_gen_derive::gen_stub_pymethods)]
127/// #[cfg_attr(not(feature = "stub-gen"), pyo3_stub_gen_derive::remove_gen_stub)]
128/// #[pymethods]
129/// impl A {
130///     #[gen_stub(override_return_type(type_repr="typing_extensions.Self", imports=("typing_extensions")))]
131///     #[new]
132///     pub fn new() -> Self {
133///         Self::default()
134///     }
135/// }
136/// #[cfg(feature = "stub-gen")]
137/// define_stub_info_gatherer!(stub_info);
138/// ```
139/// With Cargo.toml:
140/// ```toml
141/// [features]
142/// stub-gen = ["dep:pyo3-stub-gen"]
143/// [dependencies]
144/// pyo3-stub-gen = {version = "*", optional = true}
145/// pyo3-stub-gen-derive = "*"
146/// ```
147#[proc_macro_attribute]
148pub fn remove_gen_stub(_attr: TokenStream, item: TokenStream) -> TokenStream {
149    gen_stub::prune_gen_stub(item.into())
150        .unwrap_or_else(|err| err.to_compile_error())
151        .into()
152}