pure/
lib.rs

1#[cfg_attr(target_os = "macos", doc = include_str!("../../../README.md"))]
2mod readme {}
3
4use ahash::RandomState;
5use pyo3::{exceptions::PyRuntimeError, prelude::*, types::*};
6use pyo3_stub_gen::{create_exception, define_stub_info_gatherer, derive::*, module_variable};
7use std::{collections::HashMap, path::PathBuf};
8
9/// Returns the sum of two numbers as a string.
10#[gen_stub_pyfunction]
11#[pyfunction]
12fn sum(v: Vec<u32>) -> u32 {
13    v.iter().sum()
14}
15
16#[gen_stub_pyfunction]
17#[pyfunction]
18fn read_dict(dict: HashMap<usize, HashMap<usize, usize>>) {
19    for (k, v) in dict {
20        for (k2, v2) in v {
21            println!("{} {} {}", k, k2, v2);
22        }
23    }
24}
25
26#[gen_stub_pyfunction]
27#[pyfunction]
28fn create_dict(n: usize) -> HashMap<usize, Vec<usize>> {
29    let mut dict = HashMap::new();
30    for i in 0..n {
31        dict.insert(i, (0..i).collect());
32    }
33    dict
34}
35
36#[gen_stub_pyclass]
37#[pyclass(extends=PyDate)]
38struct MyDate;
39
40#[gen_stub_pyclass]
41#[pyclass(subclass)]
42#[derive(Debug)]
43struct A {
44    #[pyo3(get, set)]
45    x: usize,
46}
47
48#[gen_stub_pymethods]
49#[pymethods]
50impl A {
51    /// This is a constructor of :class:`A`.
52    #[new]
53    fn new(x: usize) -> Self {
54        Self { x }
55    }
56
57    fn show_x(&self) {
58        println!("x = {}", self.x);
59    }
60
61    fn ref_test<'a>(&self, x: Bound<'a, PyDict>) -> Bound<'a, PyDict> {
62        x
63    }
64}
65
66#[gen_stub_pyfunction]
67#[pyfunction]
68#[pyo3(signature = (x = 2))]
69fn create_a(x: usize) -> A {
70    A { x }
71}
72
73#[gen_stub_pyclass]
74#[pyclass(extends=A)]
75#[derive(Debug)]
76struct B;
77
78/// `C` only impl `FromPyObject`
79#[derive(Debug)]
80struct C {
81    x: usize,
82}
83#[gen_stub_pyfunction]
84#[pyfunction(signature = (c=None))]
85fn print_c(c: Option<C>) {
86    if let Some(c) = c {
87        println!("{}", c.x);
88    } else {
89        println!("None");
90    }
91}
92impl FromPyObject<'_> for C {
93    fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult<Self> {
94        Ok(C { x: ob.extract()? })
95    }
96}
97impl pyo3_stub_gen::PyStubType for C {
98    fn type_output() -> pyo3_stub_gen::TypeInfo {
99        usize::type_output()
100    }
101}
102
103create_exception!(pure, MyError, PyRuntimeError);
104
105/// Returns the length of the string.
106#[gen_stub_pyfunction]
107#[pyfunction]
108fn str_len(x: &str) -> PyResult<usize> {
109    Ok(x.len())
110}
111
112#[gen_stub_pyfunction]
113#[pyfunction]
114fn echo_path(path: PathBuf) -> PyResult<PathBuf> {
115    Ok(path)
116}
117
118#[gen_stub_pyfunction]
119#[pyfunction]
120fn ahash_dict() -> HashMap<String, i32, RandomState> {
121    let mut map: HashMap<String, i32, RandomState> = HashMap::with_hasher(RandomState::new());
122    map.insert("apple".to_string(), 3);
123    map.insert("banana".to_string(), 2);
124    map.insert("orange".to_string(), 5);
125    map
126}
127
128#[gen_stub_pyclass_enum]
129#[pyclass(eq, eq_int)]
130#[derive(Debug, Clone, PartialEq, Eq, Hash)]
131pub enum Number {
132    #[pyo3(name = "FLOAT")]
133    Float,
134    #[pyo3(name = "INTEGER")]
135    Integer,
136}
137
138#[gen_stub_pyclass_enum]
139#[pyclass(eq, eq_int)]
140#[pyo3(rename_all = "UPPERCASE")]
141#[derive(Debug, Clone, PartialEq, Eq, Hash)]
142pub enum NumberRenameAll {
143    /// Float variant
144    Float,
145    Integer,
146}
147
148#[gen_stub_pymethods]
149#[pymethods]
150impl Number {
151    #[getter]
152    /// Whether the number is a float.
153    fn is_float(&self) -> bool {
154        matches!(self, Self::Float)
155    }
156
157    #[getter]
158    /// Whether the number is an integer.
159    fn is_integer(&self) -> bool {
160        matches!(self, Self::Integer)
161    }
162}
163
164module_variable!("pure", "MY_CONSTANT", usize);
165
166// Test if non-any PyObject Target can be a default value
167#[gen_stub_pyfunction]
168#[pyfunction]
169#[pyo3(signature = (num = Number::Float))]
170fn default_value(num: Number) -> Number {
171    num
172}
173
174/// Initializes the Python module
175#[pymodule]
176fn pure(m: &Bound<PyModule>) -> PyResult<()> {
177    m.add("MyError", m.py().get_type::<MyError>())?;
178    m.add("MY_CONSTANT", 19937)?;
179    m.add_class::<A>()?;
180    m.add_class::<B>()?;
181    m.add_class::<MyDate>()?;
182    m.add_class::<Number>()?;
183    m.add_class::<NumberRenameAll>()?;
184    m.add_function(wrap_pyfunction!(sum, m)?)?;
185    m.add_function(wrap_pyfunction!(create_dict, m)?)?;
186    m.add_function(wrap_pyfunction!(read_dict, m)?)?;
187    m.add_function(wrap_pyfunction!(create_a, m)?)?;
188    m.add_function(wrap_pyfunction!(print_c, m)?)?;
189    m.add_function(wrap_pyfunction!(str_len, m)?)?;
190    m.add_function(wrap_pyfunction!(echo_path, m)?)?;
191    m.add_function(wrap_pyfunction!(ahash_dict, m)?)?;
192    m.add_function(wrap_pyfunction!(default_value, m)?)?;
193    Ok(())
194}
195
196define_stub_info_gatherer!(stub_info);
197
198/// Test of unit test for testing link problem
199#[cfg(test)]
200mod test {
201    #[test]
202    fn test() {
203        assert_eq!(2 + 2, 4);
204    }
205}
pure/lib.rs

pure/
lib.rs
