pyo3_stub_gen/stub_type/
builtins.rs

1//! Define PyStubType for built-in types based on <https://pyo3.rs/v0.22.2/conversions/tables#argument-types>
2
3use crate::runtime::{union_type, PyRuntimeType};
4use crate::stub_type::*;
5use ::pyo3::prelude::*;
6use ::pyo3::types::{PyBool, PyComplex, PyFloat, PyInt, PyString};
7use std::{
8    borrow::Cow,
9    ffi::{OsStr, OsString},
10    path::PathBuf,
11    rc::Rc,
12    sync::Arc,
13    time::SystemTime,
14};
15
16macro_rules! impl_builtin {
17    ($ty:ty, $pytype:expr, $py_type_obj:ty) => {
18        impl PyStubType for $ty {
19            fn type_output() -> TypeInfo {
20                TypeInfo::builtin($pytype)
21            }
22        }
23        impl PyRuntimeType for $ty {
24            fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
25                Ok(py.get_type::<$py_type_obj>().into_any())
26            }
27        }
28    };
29}
30
31macro_rules! impl_with_module {
32    ($ty:ty, $pytype:expr, $module:expr) => {
33        impl PyStubType for $ty {
34            fn type_output() -> TypeInfo {
35                TypeInfo::with_module($pytype, $module.into())
36            }
37        }
38        impl PyRuntimeType for $ty {
39            fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
40                // Import the type from the Python module
41                let module = py.import($module)?;
42                // Extract just the type name from "module.TypeName"
43                let type_name = $pytype.rsplit('.').next().unwrap_or($pytype);
44                module.getattr(type_name)
45            }
46        }
47    };
48}
49
50// NOTE:
51impl PyStubType for () {
52    fn type_output() -> TypeInfo {
53        TypeInfo::none()
54    }
55}
56impl PyRuntimeType for () {
57    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
58        // None type in Python
59        Ok(py.get_type::<::pyo3::types::PyNone>().into_any())
60    }
61}
62
63impl_builtin!(bool, "bool", PyBool);
64impl_builtin!(u8, "int", PyInt);
65impl_builtin!(u16, "int", PyInt);
66impl_builtin!(u32, "int", PyInt);
67impl_builtin!(u64, "int", PyInt);
68impl_builtin!(u128, "int", PyInt);
69impl_builtin!(usize, "int", PyInt);
70impl_builtin!(i8, "int", PyInt);
71impl_builtin!(i16, "int", PyInt);
72impl_builtin!(i32, "int", PyInt);
73impl_builtin!(i64, "int", PyInt);
74impl_builtin!(i128, "int", PyInt);
75impl_builtin!(isize, "int", PyInt);
76impl_builtin!(f32, "float", PyFloat);
77impl_builtin!(f64, "float", PyFloat);
78impl_builtin!(num_complex::Complex32, "complex", PyComplex);
79impl_builtin!(num_complex::Complex64, "complex", PyComplex);
80
81impl_builtin!(char, "str", PyString);
82impl_builtin!(&str, "str", PyString);
83impl_builtin!(OsStr, "str", PyString);
84impl_builtin!(String, "str", PyString);
85impl_builtin!(OsString, "str", PyString);
86impl_builtin!(Cow<'_, str>, "str", PyString);
87impl_builtin!(Cow<'_, OsStr>, "str", PyString);
88impl_builtin!(Cow<'_, [u8]>, "bytes", ::pyo3::types::PyBytes);
89
90#[cfg(feature = "ordered-float")]
91mod impl_ordered_float {
92    use super::*;
93    use ::pyo3::types::PyFloat;
94    impl_builtin!(ordered_float::NotNan<f32>, "float", PyFloat);
95    impl_builtin!(ordered_float::NotNan<f64>, "float", PyFloat);
96    impl_builtin!(ordered_float::OrderedFloat<f32>, "float", PyFloat);
97    impl_builtin!(ordered_float::OrderedFloat<f64>, "float", PyFloat);
98}
99
100impl PyStubType for PathBuf {
101    fn type_output() -> TypeInfo {
102        TypeInfo::with_module("pathlib.Path", "pathlib".into())
103    }
104    fn type_input() -> TypeInfo {
105        TypeInfo::builtin("str")
106            | TypeInfo::with_module("os.PathLike", "os".into())
107            | TypeInfo::with_module("pathlib.Path", "pathlib".into())
108    }
109}
110impl PyRuntimeType for PathBuf {
111    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
112        let pathlib = py.import("pathlib")?;
113        pathlib.getattr("Path")
114    }
115}
116
117impl<Tz: chrono::TimeZone> PyStubType for chrono::DateTime<Tz> {
118    fn type_output() -> TypeInfo {
119        TypeInfo::with_module("datetime.datetime", "datetime".into())
120    }
121}
122impl<Tz: chrono::TimeZone> PyRuntimeType for chrono::DateTime<Tz> {
123    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
124        let datetime = py.import("datetime")?;
125        datetime.getattr("datetime")
126    }
127}
128
129impl_with_module!(SystemTime, "datetime.datetime", "datetime");
130impl_with_module!(chrono::NaiveDateTime, "datetime.datetime", "datetime");
131impl_with_module!(chrono::NaiveDate, "datetime.date", "datetime");
132impl_with_module!(chrono::NaiveTime, "datetime.time", "datetime");
133impl_with_module!(chrono::FixedOffset, "datetime.tzinfo", "datetime");
134impl_with_module!(chrono::Utc, "datetime.tzinfo", "datetime");
135impl_with_module!(std::time::Duration, "datetime.timedelta", "datetime");
136impl_with_module!(chrono::Duration, "datetime.timedelta", "datetime");
137impl_with_module!(time::Duration, "datetime.timedelta", "datetime");
138impl_with_module!(time::Date, "datetime.date", "datetime");
139impl_with_module!(time::OffsetDateTime, "datetime.datetime", "datetime");
140impl_with_module!(time::PrimitiveDateTime, "datetime.datetime", "datetime");
141impl_with_module!(time::UtcDateTime, "datetime.datetime", "datetime");
142impl_with_module!(time::Time, "datetime.time", "datetime");
143impl_with_module!(time::UtcOffset, "datetime.tzinfo", "datetime");
144impl_with_module!(std::net::Ipv4Addr, "ipaddress.IPv4Address", "ipaddress");
145impl_with_module!(std::net::Ipv6Addr, "ipaddress.IPv6Address", "ipaddress");
146
147impl PyStubType for std::net::IpAddr {
148    fn type_output() -> TypeInfo {
149        TypeInfo::with_module(
150            "ipaddress.IPv4Address",
151            ModuleRef::Named("ipaddress".to_string()),
152        ) | TypeInfo::with_module(
153            "ipaddress.IPv6Address",
154            ModuleRef::Named("ipaddress".to_string()),
155        )
156    }
157}
158impl PyRuntimeType for std::net::IpAddr {
159    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
160        let module = py.import("ipaddress")?;
161        let type_v4 = module.getattr("IPv4Address")?;
162        let type_v6 = module.getattr("IPv6Address")?;
163        union_type(py, &[type_v4, type_v6])
164    }
165}
166
167impl<T: PyStubType> PyStubType for &T {
168    fn type_input() -> TypeInfo {
169        T::type_input()
170    }
171    fn type_output() -> TypeInfo {
172        T::type_output()
173    }
174}
175impl<T: PyRuntimeType> PyRuntimeType for &T {
176    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
177        T::runtime_type_object(py)
178    }
179}
180
181impl<T: PyStubType> PyStubType for Rc<T> {
182    fn type_input() -> TypeInfo {
183        T::type_input()
184    }
185    fn type_output() -> TypeInfo {
186        T::type_output()
187    }
188}
189impl<T: PyRuntimeType> PyRuntimeType for Rc<T> {
190    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
191        T::runtime_type_object(py)
192    }
193}
194
195impl<T: PyStubType> PyStubType for Arc<T> {
196    fn type_input() -> TypeInfo {
197        T::type_input()
198    }
199    fn type_output() -> TypeInfo {
200        T::type_output()
201    }
202}
203impl<T: PyRuntimeType> PyRuntimeType for Arc<T> {
204    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
205        T::runtime_type_object(py)
206    }
207}
pyo3_stub_gen/stub_type/builtins.rs

pyo3_stub_gen/stub_type/
builtins.rs
