Skip to content

Commit d09f8e2

Browse files
TptTpt
committed
Introspection: implement output type
- Adds a new 'returns' field in the introspection data that contains the returned type annotation - Generate it using a new IntoPyObject::OUTPUT_TYPE associated constant (similar to INPUT_TYPE) - Add the internal PyReturnType trait to allow getting easily T::RETURN_TYPE when the function returns Result<T, E> - Take care of the special case of `()` converted to `None` and not `()` - Adding `OUTPUT_TYPE` to the `IntoPyObject` will be a follow-up (to avoid a huge PR)
1 parent 8fa2d60 commit d09f8e2

File tree

16 files changed

+176
-36
lines changed

16 files changed

+176
-36
lines changed

‎newsfragments/5208.added.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Introspection and sub generation: add basic return type support

‎pyo3-introspection/src/introspection.rs

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,8 @@ fn convert_members(
121121
arguments,
122122
parent: _,
123123
decorators,
124-
} => functions.push(convert_function(name, arguments, decorators)),
124+
returns,
125+
} => functions.push(convert_function(name, arguments, decorators, returns)),
125126
}
126127
}
127128
Ok((modules, classes, functions))
@@ -170,7 +171,12 @@ fn convert_class(
170171
})
171172
}
172173

173-
fn convert_function(name: &str, arguments: &ChunkArguments, decorators: &[String]) -> Function {
174+
fn convert_function(
175+
name: &str,
176+
arguments: &ChunkArguments,
177+
decorators: &[String],
178+
returns: &Option<String>,
179+
) -> Function {
174180
Function {
175181
name: name.into(),
176182
decorators: decorators.to_vec(),
@@ -187,6 +193,7 @@ fn convert_function(name: &str, arguments: &ChunkArguments, decorators: &[String
187193
.as_ref()
188194
.map(convert_variable_length_argument),
189195
},
196+
returns: returns.clone(),
190197
}
191198
}
192199

@@ -382,6 +389,8 @@ enum Chunk {
382389
parent: Option<String>,
383390
#[serde(default)]
384391
decorators: Vec<String>,
392+
#[serde(default)]
393+
returns: Option<String>,
385394
},
386395
}
387396

‎pyo3-introspection/src/model.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ pub struct Function {
1919
/// decorator like 'property' or 'staticmethod'
2020
pub decorators: Vec<String>,
2121
pub arguments: Arguments,
22+
/// return type
23+
pub returns: Option<String>,
2224
}
2325

2426
#[derive(Debug, Eq, PartialEq, Clone, Hash)]

‎pyo3-introspection/src/stubs.rs

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -108,17 +108,22 @@ fn function_stubs(function: &Function, modules_to_import: &mut BTreeSet<String>)
108108
if let Some(argument) = &function.arguments.kwarg {
109109
parameters.push(format!("**{}", variable_length_argument_stub(argument)));
110110
}
111-
let output = format!("def {}({}): ...", function.name, parameters.join(", "));
112-
if function.decorators.is_empty() {
113-
return output;
114-
}
115111
let mut buffer = String::new();
116112
for decorator in &function.decorators {
117113
buffer.push('@');
118114
buffer.push_str(decorator);
119115
buffer.push('\n');
120116
}
121-
buffer.push_str(&output);
117+
buffer.push_str("def ");
118+
buffer.push_str(&function.name);
119+
buffer.push('(');
120+
buffer.push_str(&parameters.join(", "));
121+
buffer.push(')');
122+
if let Some(returns) = &function.returns {
123+
buffer.push_str(" -> ");
124+
buffer.push_str(annotation_stub(returns, modules_to_import));
125+
}
126+
buffer.push_str(": ...");
122127
buffer
123128
}
124129

@@ -132,11 +137,7 @@ fn argument_stub(argument: &Argument, modules_to_import: &mut BTreeSet<String>)
132137
let mut output = argument.name.clone();
133138
if let Some(annotation) = &argument.annotation {
134139
output.push_str(": ");
135-
output.push_str(annotation);
136-
if let Some((module, _)) = annotation.rsplit_once('.') {
137-
// TODO: this is very naive
138-
modules_to_import.insert(module.into());
139-
}
140+
output.push_str(annotation_stub(annotation, modules_to_import));
140141
}
141142
if let Some(default_value) = &argument.default_value {
142143
output.push_str(if argument.annotation.is_some() {
@@ -153,6 +154,14 @@ fn variable_length_argument_stub(argument: &VariableLengthArgument) -> String {
153154
argument.name.clone()
154155
}
155156

157+
fn annotation_stub<'a>(annotation: &'a str, modules_to_import: &mut BTreeSet<String>) -> &'a str {
158+
if let Some((module, _)) = annotation.rsplit_once('.') {
159+
// TODO: this is very naive
160+
modules_to_import.insert(module.into());
161+
}
162+
annotation
163+
}
164+
156165
#[cfg(test)]
157166
mod tests {
158167
use super::*;
@@ -186,9 +195,10 @@ mod tests {
186195
name: "kwarg".into(),
187196
}),
188197
},
198+
returns: Some("list[str]".into()),
189199
};
190200
assert_eq!(
191-
"def func(posonly, /, arg, *varargs, karg: str, **kwarg): ...",
201+
"def func(posonly, /, arg, *varargs, karg: str, **kwarg) -> list[str]: ...",
192202
function_stubs(&function, &mut BTreeSet::new())
193203
)
194204
}
@@ -217,6 +227,7 @@ mod tests {
217227
}],
218228
kwarg: None,
219229
},
230+
returns: None,
220231
};
221232
assert_eq!(
222233
"def afunc(posonly=1, /, arg=True, *, karg: str = \"foo\"): ...",

‎pyo3-macros-backend/src/introspection.rs

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use std::mem::take;
2121
use std::sync::atomic::{AtomicUsize, Ordering};
2222
use syn::ext::IdentExt;
2323
use syn::visit_mut::{visit_type_mut, VisitMut};
24-
use syn::{Attribute, Ident, Type, TypePath};
24+
use syn::{Attribute, Ident, ReturnType, Type, TypePath};
2525

2626
static GLOBAL_COUNTER_FOR_UNIQUE_NAMES: AtomicUsize = AtomicUsize::new(0);
2727

@@ -99,12 +99,14 @@ pub fn class_introspection_code(
9999
.emit(pyo3_crate_path)
100100
}
101101

102+
#[allow(clippy::too_many_arguments)]
102103
pub fn function_introspection_code(
103104
pyo3_crate_path: &PyO3CratePath,
104105
ident: Option<&Ident>,
105106
name: &str,
106107
signature: &FunctionSignature<'_>,
107108
first_argument: Option<&'static str>,
109+
returns: ReturnType,
108110
decorators: impl IntoIterator<Item = String>,
109111
parent: Option<&Type>,
110112
) -> TokenStream {
@@ -115,6 +117,25 @@ pub fn function_introspection_code(
115117
"arguments",
116118
arguments_introspection_data(signature, first_argument, parent),
117119
),
120+
(
121+
"returns",
122+
match returns {
123+
ReturnType::Default => IntrospectionNode::String("None".into()),
124+
ReturnType::Type(_, ty) => match *ty {
125+
Type::Tuple(t) if t.elems.is_empty() => {
126+
// () is converted to None in return types
127+
IntrospectionNode::String("None".into())
128+
}
129+
mut ty => {
130+
if let Some(class_type) = parent {
131+
replace_self(&mut ty, class_type);
132+
}
133+
ty = ty.elide_lifetimes();
134+
IntrospectionNode::OutputType { rust_type: ty }
135+
}
136+
},
137+
},
138+
),
118139
]);
119140
if let Some(ident) = ident {
120141
desc.insert(
@@ -290,6 +311,7 @@ enum IntrospectionNode<'a> {
290311
String(Cow<'a, str>),
291312
IntrospectionId(Option<Cow<'a, Type>>),
292313
InputType { rust_type: Type, nullable: bool },
314+
OutputType { rust_type: Type },
293315
Map(HashMap<&'static str, IntrospectionNode<'a>>),
294316
List(Vec<IntrospectionNode<'a>>),
295317
}
@@ -340,6 +362,11 @@ impl IntrospectionNode<'_> {
340362
}
341363
content.push_str("\"");
342364
}
365+
Self::OutputType { rust_type } => {
366+
content.push_str("\"");
367+
content.push_tokens(quote! { <#rust_type as #pyo3_crate_path::impl_::introspection::PyReturnType>::OUTPUT_TYPE });
368+
content.push_str("\"");
369+
}
343370
Self::Map(map) => {
344371
content.push_str("{");
345372
for (i, (key, value)) in map.into_iter().enumerate() {

‎pyo3-macros-backend/src/method.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -445,6 +445,8 @@ pub struct FnSpec<'a> {
445445
pub asyncness: Option<syn::Token![async]>,
446446
pub unsafety: Option<syn::Token![unsafe]>,
447447
pub warnings: Vec<PyFunctionWarning>,
448+
#[cfg(feature = "experimental-inspect")]
449+
pub output: syn::ReturnType,
448450
}
449451

450452
pub fn parse_method_receiver(arg: &syn::FnArg) -> Result<SelfType> {
@@ -526,6 +528,8 @@ impl<'a> FnSpec<'a> {
526528
asyncness: sig.asyncness,
527529
unsafety: sig.unsafety,
528530
warnings,
531+
#[cfg(feature = "experimental-inspect")]
532+
output: sig.output.clone(),
529533
})
530534
}
531535

‎pyo3-macros-backend/src/pyclass.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1670,6 +1670,8 @@ fn complex_enum_struct_variant_new<'a>(
16701670
asyncness: None,
16711671
unsafety: None,
16721672
warnings: vec![],
1673+
#[cfg(feature = "experimental-inspect")]
1674+
output: syn::ReturnType::Default,
16731675
};
16741676

16751677
crate::pymethod::impl_py_method_def_new(&variant_cls_type, &spec, ctx)
@@ -1725,6 +1727,8 @@ fn complex_enum_tuple_variant_new<'a>(
17251727
asyncness: None,
17261728
unsafety: None,
17271729
warnings: vec![],
1730+
#[cfg(feature = "experimental-inspect")]
1731+
output: syn::ReturnType::Default,
17281732
};
17291733

17301734
crate::pymethod::impl_py_method_def_new(&variant_cls_type, &spec, ctx)
@@ -1750,6 +1754,8 @@ fn complex_enum_variant_field_getter<'a>(
17501754
asyncness: None,
17511755
unsafety: None,
17521756
warnings: vec![],
1757+
#[cfg(feature = "experimental-inspect")]
1758+
output: syn::ReturnType::Type(Token![->](field_span), Box::new(variant_cls_type.clone())),
17531759
};
17541760

17551761
let property_type = crate::pymethod::PropertyType::Function {

‎pyo3-macros-backend/src/pyfunction.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,7 @@ pub fn impl_wrap_pyfunction(
390390
&name.to_string(),
391391
&signature,
392392
None,
393+
func.sig.output.clone(),
393394
[] as [String; 0],
394395
None,
395396
);
@@ -410,6 +411,8 @@ pub fn impl_wrap_pyfunction(
410411
asyncness: func.sig.asyncness,
411412
unsafety: func.sig.unsafety,
412413
warnings,
414+
#[cfg(feature = "experimental-inspect")]
415+
output: func.sig.output.clone(),
413416
};
414417

415418
let wrapper_ident = format_ident!("__pyfunction_{}", spec.name);

‎pyo3-macros-backend/src/pyimpl.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,6 +367,7 @@ fn method_introspection_code(spec: &FnSpec<'_>, parent: &syn::Type, ctx: &Ctx) -
367367

368368
// We introduce self/cls argument and setup decorators
369369
let mut first_argument = None;
370+
let mut output = spec.output.clone();
370371
let mut decorators = Vec::new();
371372
match &spec.tp {
372373
FnType::Getter(_) => {
@@ -382,6 +383,7 @@ fn method_introspection_code(spec: &FnSpec<'_>, parent: &syn::Type, ctx: &Ctx) -
382383
}
383384
FnType::FnNew | FnType::FnNewClass(_) => {
384385
first_argument = Some("cls");
386+
output = syn::ReturnType::Default; // The __new__ Python function return type is None
385387
}
386388
FnType::FnClass(_) => {
387389
first_argument = Some("cls");
@@ -404,6 +406,7 @@ fn method_introspection_code(spec: &FnSpec<'_>, parent: &syn::Type, ctx: &Ctx) -
404406
&name,
405407
&spec.signature,
406408
first_argument,
409+
output,
407410
decorators,
408411
Some(parent),
409412
)

‎pytests/stubs/pyclasses.pyi

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,33 @@
11
import typing
22

33
class AssertingBaseClass:
4-
def __new__(cls, /, expected_type: typing.Any): ...
4+
def __new__(cls, /, expected_type: typing.Any) -> None: ...
55

66
class ClassWithDecorators:
7-
def __new__(cls, /): ...
7+
def __new__(cls, /) -> None: ...
88
@property
9-
def attr(self, /): ...
9+
def attr(self, /) -> int: ...
1010
@attr.setter
11-
def attr(self, /, value: int): ...
11+
def attr(self, /, value: int) -> None: ...
1212
@classmethod
1313
@property
14-
def cls_attribute(cls, /): ...
14+
def cls_attribute(cls, /) -> int: ...
1515
@classmethod
16-
def cls_method(cls, /): ...
16+
def cls_method(cls, /) -> int: ...
1717
@staticmethod
18-
def static_method(): ...
18+
def static_method() -> int: ...
1919

2020
class ClassWithoutConstructor: ...
2121

2222
class EmptyClass:
23-
def __len__(self, /): ...
24-
def __new__(cls, /): ...
25-
def method(self, /): ...
23+
def __len__(self, /) -> int: ...
24+
def __new__(cls, /) -> None: ...
25+
def method(self, /) -> None: ...
2626

2727
class PyClassIter:
28-
def __new__(cls, /): ...
29-
def __next__(self, /): ...
28+
def __new__(cls, /) -> None: ...
29+
def __next__(self, /) -> int: ...
3030

3131
class PyClassThreadIter:
32-
def __new__(cls, /): ...
33-
def __next__(self, /): ...
32+
def __new__(cls, /) -> None: ...
33+
def __next__(self, /) -> int: ...

‎pytests/stubs/pyfunctions.pyi

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,24 @@
11
import typing
22

3-
def args_kwargs(*args, **kwargs): ...
4-
def none(): ...
5-
def positional_only(a: typing.Any, /, b: typing.Any): ...
3+
def args_kwargs(*args, **kwargs) -> typing.Any: ...
4+
def none() -> None: ...
5+
def positional_only(a: typing.Any, /, b: typing.Any) -> typing.Any: ...
66
def simple(
77
a: typing.Any, b: typing.Any | None = None, *, c: typing.Any | None = None
8-
): ...
8+
) -> typing.Any: ...
99
def simple_args(
1010
a: typing.Any, b: typing.Any | None = None, *args, c: typing.Any | None = None
11-
): ...
11+
) -> typing.Any: ...
1212
def simple_args_kwargs(
1313
a: typing.Any,
1414
b: typing.Any | None = None,
1515
*args,
1616
c: typing.Any | None = None,
1717
**kwargs,
18-
): ...
18+
) -> typing.Any: ...
1919
def simple_kwargs(
2020
a: typing.Any, b: typing.Any | None = None, c: typing.Any | None = None, **kwargs
21-
): ...
22-
def with_typed_args(a: bool = False, b: int = 0, c: float = 0.0, d: str = ""): ...
21+
) -> typing.Any: ...
22+
def with_typed_args(
23+
a: bool = False, b: int = 0, c: float = 0.0, d: str = ""
24+
) -> typing.Any: ...

‎src/conversion.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,16 @@ pub trait IntoPyObject<'py>: Sized {
4747
/// The type returned in the event of a conversion error.
4848
type Error: Into<PyErr>;
4949

50+
/// Extracts the type hint information for this type when it appears as a return value.
51+
///
52+
/// For example, `Vec<u32>` would return `List[int]`.
53+
/// The default implementation returns `Any`, which is correct for any type.
54+
///
55+
/// For most types, the return value for this method will be identical to that of [`FromPyObject::INPUT_TYPE`].
56+
/// It may be different for some types, such as `Dict`, to allow duck-typing: functions return `Dict` but take `Mapping` as argument.
57+
#[cfg(feature = "experimental-inspect")]
58+
const OUTPUT_TYPE: &'static str = "typing.Any";
59+
5060
/// Performs the conversion.
5161
fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error>;
5262

0 commit comments

Comments
 (0)