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 = "jiff-02")]
91mod impl_jiff_02 {
92    use super::*;
93
94    impl_with_module!(jiff_02::Timestamp, "datetime.datetime", "datetime");
95    impl_with_module!(jiff_02::Zoned, "datetime.datetime", "datetime");
96    impl_with_module!(jiff_02::civil::DateTime, "datetime.datetime", "datetime");
97    impl_with_module!(jiff_02::civil::Date, "datetime.date", "datetime");
98    impl_with_module!(jiff_02::civil::Time, "datetime.time", "datetime");
99    impl_with_module!(jiff_02::civil::ISOWeekDate, "datetime.date", "datetime");
100    impl_with_module!(jiff_02::tz::Offset, "datetime.tzinfo", "datetime");
101    impl_with_module!(jiff_02::tz::TimeZone, "datetime.tzinfo", "datetime");
102    impl_with_module!(jiff_02::SignedDuration, "datetime.timedelta", "datetime");
103    impl_with_module!(jiff_02::Span, "datetime.timedelta", "datetime");
104}
105
106#[cfg(feature = "ordered-float")]
107mod impl_ordered_float {
108    use super::*;
109    use ::pyo3::types::PyFloat;
110    impl_builtin!(ordered_float::NotNan<f32>, "float", PyFloat);
111    impl_builtin!(ordered_float::NotNan<f64>, "float", PyFloat);
112    impl_builtin!(ordered_float::OrderedFloat<f32>, "float", PyFloat);
113    impl_builtin!(ordered_float::OrderedFloat<f64>, "float", PyFloat);
114}
115
116impl PyStubType for PathBuf {
117    fn type_output() -> TypeInfo {
118        TypeInfo::with_module("pathlib.Path", "pathlib".into())
119    }
120    fn type_input() -> TypeInfo {
121        TypeInfo::builtin("str")
122            | TypeInfo::with_module("os.PathLike", "os".into())
123            | TypeInfo::with_module("pathlib.Path", "pathlib".into())
124    }
125}
126impl PyRuntimeType for PathBuf {
127    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
128        let pathlib = py.import("pathlib")?;
129        pathlib.getattr("Path")
130    }
131}
132
133impl<Tz: chrono::TimeZone> PyStubType for chrono::DateTime<Tz> {
134    fn type_output() -> TypeInfo {
135        TypeInfo::with_module("datetime.datetime", "datetime".into())
136    }
137}
138impl<Tz: chrono::TimeZone> PyRuntimeType for chrono::DateTime<Tz> {
139    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
140        let datetime = py.import("datetime")?;
141        datetime.getattr("datetime")
142    }
143}
144
145impl_with_module!(SystemTime, "datetime.datetime", "datetime");
146impl_with_module!(chrono::NaiveDateTime, "datetime.datetime", "datetime");
147impl_with_module!(chrono::NaiveDate, "datetime.date", "datetime");
148impl_with_module!(chrono::NaiveTime, "datetime.time", "datetime");
149impl_with_module!(chrono::FixedOffset, "datetime.tzinfo", "datetime");
150impl_with_module!(chrono::Utc, "datetime.tzinfo", "datetime");
151impl_with_module!(std::time::Duration, "datetime.timedelta", "datetime");
152impl_with_module!(chrono::Duration, "datetime.timedelta", "datetime");
153impl_with_module!(time::Duration, "datetime.timedelta", "datetime");
154impl_with_module!(time::Date, "datetime.date", "datetime");
155impl_with_module!(time::OffsetDateTime, "datetime.datetime", "datetime");
156impl_with_module!(time::PrimitiveDateTime, "datetime.datetime", "datetime");
157impl_with_module!(time::UtcDateTime, "datetime.datetime", "datetime");
158impl_with_module!(time::Time, "datetime.time", "datetime");
159impl_with_module!(time::UtcOffset, "datetime.tzinfo", "datetime");
160impl_with_module!(std::net::Ipv4Addr, "ipaddress.IPv4Address", "ipaddress");
161impl_with_module!(std::net::Ipv6Addr, "ipaddress.IPv6Address", "ipaddress");
162
163impl PyStubType for std::net::IpAddr {
164    fn type_output() -> TypeInfo {
165        TypeInfo::with_module(
166            "ipaddress.IPv4Address",
167            ModuleRef::Named("ipaddress".to_string()),
168        ) | TypeInfo::with_module(
169            "ipaddress.IPv6Address",
170            ModuleRef::Named("ipaddress".to_string()),
171        )
172    }
173}
174impl PyRuntimeType for std::net::IpAddr {
175    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
176        let module = py.import("ipaddress")?;
177        let type_v4 = module.getattr("IPv4Address")?;
178        let type_v6 = module.getattr("IPv6Address")?;
179        union_type(py, &[type_v4, type_v6])
180    }
181}
182
183#[cfg(feature = "uuid")]
184impl PyStubType for uuid::Uuid {
185    fn type_output() -> TypeInfo {
186        TypeInfo::with_module("uuid.UUID", "uuid".into())
187    }
188}
189
190impl<T: PyStubType> PyStubType for &T {
191    fn type_input() -> TypeInfo {
192        T::type_input()
193    }
194    fn type_output() -> TypeInfo {
195        T::type_output()
196    }
197}
198impl<T: PyRuntimeType> PyRuntimeType for &T {
199    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
200        T::runtime_type_object(py)
201    }
202}
203
204impl<T: PyStubType> PyStubType for Rc<T> {
205    fn type_input() -> TypeInfo {
206        T::type_input()
207    }
208    fn type_output() -> TypeInfo {
209        T::type_output()
210    }
211}
212impl<T: PyRuntimeType> PyRuntimeType for Rc<T> {
213    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
214        T::runtime_type_object(py)
215    }
216}
217
218impl<T: PyStubType> PyStubType for Arc<T> {
219    fn type_input() -> TypeInfo {
220        T::type_input()
221    }
222    fn type_output() -> TypeInfo {
223        T::type_output()
224    }
225}
226impl<T: PyRuntimeType> PyRuntimeType for Arc<T> {
227    fn runtime_type_object(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
228        T::runtime_type_object(py)
229    }
230}
pyo3_stub_gen/stub_type/builtins.rs

pyo3_stub_gen/stub_type/
builtins.rs
