Skip to content

Commit 1faf389

Browse files
TptTpt
authored
TypeHint: distinguish local and builtins (#5554)
* TypeHint: distinguish local and builtins - local: unknown module - builtins: Python builtins This distinction allows to support overridden Python builtins, in this case the stub import them using "from builtins impo X as X2" * Remove str_equals
1 parent 7d9464b commit 1faf389

File tree

6 files changed

+90
-34
lines changed

6 files changed

+90
-34
lines changed

‎pyo3-introspection/src/introspection.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,7 @@ fn convert_type_hint(arg: &ChunkTypeHint) -> TypeHint {
253253

254254
fn convert_type_hint_expr(expr: &ChunkTypeHintExpr) -> TypeHintExpr {
255255
match expr {
256+
ChunkTypeHintExpr::Local { id } => TypeHintExpr::Local { id: id.clone() },
256257
ChunkTypeHintExpr::Builtin { id } => TypeHintExpr::Builtin { id: id.clone() },
257258
ChunkTypeHintExpr::Attribute { module, attr } => TypeHintExpr::Attribute {
258259
module: module.clone(),
@@ -469,6 +470,9 @@ enum ChunkTypeHint {
469470
#[derive(Deserialize)]
470471
#[serde(tag = "type", rename_all = "lowercase")]
471472
enum ChunkTypeHintExpr {
473+
Local {
474+
id: String,
475+
},
472476
Builtin {
473477
id: String,
474478
},

‎pyo3-introspection/src/model.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,8 @@ pub enum TypeHint {
7777
/// A type hint annotation as an AST fragment
7878
#[derive(Debug, Eq, PartialEq, Clone, Hash)]
7979
pub enum TypeHintExpr {
80+
/// A local identifier which module is unknown
81+
Local { id: String },
8082
/// A Python builtin like `int`
8183
Builtin { id: String },
8284
/// The attribute of a python object like `{value}.{attr}`

‎pyo3-introspection/src/stubs.rs‎

Lines changed: 47 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -254,7 +254,7 @@ impl Imports {
254254
.map(|c| c.name.clone())
255255
.chain(module.functions.iter().map(|f| f.name.clone()))
256256
.chain(module.attributes.iter().map(|a| a.name.clone()))
257-
.chain(elements_used_in_annotations.builtins)
257+
.chain(elements_used_in_annotations.locals)
258258
{
259259
local_name_to_module_and_attribute.insert(name.clone(), (None, name.clone()));
260260
}
@@ -324,18 +324,24 @@ impl Imports {
324324
local_name.clone(),
325325
(normalized_module.clone(), root_attr.to_owned()),
326326
);
327-
import_for_module.push(if local_name == root_attr {
328-
local_name
329-
} else {
330-
format!("{root_attr} as {local_name}")
331-
});
327+
let is_not_aliased_builtin =
328+
normalized_module.as_deref() == Some("builtins") && local_name == root_attr;
329+
if !is_not_aliased_builtin {
330+
import_for_module.push(if local_name == root_attr {
331+
local_name
332+
} else {
333+
format!("{root_attr} as {local_name}")
334+
});
335+
}
332336
}
333337
}
334338
if let Some(module) = normalized_module {
335-
imports.push(format!(
336-
"from {module} import {}",
337-
import_for_module.join(", ")
338-
));
339+
if !import_for_module.is_empty() {
340+
imports.push(format!(
341+
"from {module} import {}",
342+
import_for_module.join(", ")
343+
));
344+
}
339345
}
340346
}
341347

@@ -344,7 +350,14 @@ impl Imports {
344350

345351
fn serialize_type_hint(&self, expr: &TypeHintExpr, buffer: &mut String) {
346352
match expr {
347-
TypeHintExpr::Builtin { id } => buffer.push_str(id),
353+
TypeHintExpr::Local { id } => buffer.push_str(id),
354+
TypeHintExpr::Builtin { id } => {
355+
let alias = self
356+
.renaming
357+
.get(&("builtins".to_string(), id.clone()))
358+
.expect("All type hint attributes should have been visited");
359+
buffer.push_str(alias)
360+
}
348361
TypeHintExpr::Attribute { module, attr } => {
349362
let alias = self
350363
.renaming
@@ -379,14 +392,14 @@ impl Imports {
379392
struct ElementsUsedInAnnotations {
380393
/// module -> name
381394
module_members: BTreeMap<String, BTreeSet<String>>,
382-
builtins: BTreeSet<String>,
395+
locals: BTreeSet<String>,
383396
}
384397

385398
impl ElementsUsedInAnnotations {
386399
fn new() -> Self {
387400
Self {
388401
module_members: BTreeMap::new(),
389-
builtins: BTreeSet::new(),
402+
locals: BTreeSet::new(),
390403
}
391404
}
392405

@@ -401,7 +414,10 @@ impl ElementsUsedInAnnotations {
401414
self.walk_function(function);
402415
}
403416
if module.incomplete {
404-
self.builtins.insert("str".into());
417+
self.module_members
418+
.entry("builtins".into())
419+
.or_default()
420+
.insert("str".into());
405421
self.module_members
406422
.entry("_typeshed".into())
407423
.or_default()
@@ -426,7 +442,7 @@ impl ElementsUsedInAnnotations {
426442

427443
fn walk_function(&mut self, function: &Function) {
428444
for decorator in &function.decorators {
429-
self.builtins.insert(decorator.clone());
445+
self.locals.insert(decorator.clone()); // TODO: better decorator support
430446
}
431447
for arg in function
432448
.arguments
@@ -463,8 +479,14 @@ impl ElementsUsedInAnnotations {
463479

464480
fn walk_type_hint_expr(&mut self, expr: &TypeHintExpr) {
465481
match expr {
482+
TypeHintExpr::Local { id } => {
483+
self.locals.insert(id.clone());
484+
}
466485
TypeHintExpr::Builtin { id } => {
467-
self.builtins.insert(id.clone());
486+
self.module_members
487+
.entry("builtins".into())
488+
.or_default()
489+
.insert(id.clone());
468490
}
469491
TypeHintExpr::Attribute { module, attr } => {
470492
self.module_members
@@ -593,6 +615,9 @@ mod tests {
593615
module: "bat".into(),
594616
attr: "A".into(),
595617
},
618+
TypeHintExpr::Local { id: "int".into() },
619+
TypeHintExpr::Builtin { id: "int".into() },
620+
TypeHintExpr::Builtin { id: "float".into() },
596621
],
597622
},
598623
],
@@ -628,11 +653,15 @@ mod tests {
628653
&[
629654
"from _typeshed import Incomplete",
630655
"from bat import A as A2",
631-
"from foo import A as A3, B"
656+
"from builtins import int as int2",
657+
"from foo import A as A3, B",
632658
]
633659
);
634660
let mut output = String::new();
635661
imports.serialize_type_hint(&big_type, &mut output);
636-
assert_eq!(output, "dict[A, A | A3.C | A3.D | B | A2]");
662+
assert_eq!(
663+
output,
664+
"dict[A, A | A3.C | A3.D | B | A2 | int | int2 | float]"
665+
);
637666
}
638667
}

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

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -105,14 +105,14 @@ pub fn function_introspection_code(
105105
match returns {
106106
ReturnType::Default => IntrospectionNode::ConstantType {
107107
name: "None",
108-
module: None,
108+
module: "builtins",
109109
},
110110
ReturnType::Type(_, ty) => match *ty {
111111
Type::Tuple(t) if t.elems.is_empty() => {
112112
// () is converted to None in return types
113113
IntrospectionNode::ConstantType {
114114
name: "None",
115-
module: None,
115+
module: "builtins",
116116
}
117117
}
118118
mut ty => {
@@ -190,7 +190,7 @@ pub fn attribute_introspection_code(
190190
// typing.Final[typing.literal[value]]
191191
IntrospectionNode::ConstantType {
192192
name: "Final",
193-
module: Some("typing"),
193+
module: "typing",
194194
}
195195
} else {
196196
IntrospectionNode::OutputType {
@@ -361,7 +361,7 @@ enum IntrospectionNode<'a> {
361361
},
362362
ConstantType {
363363
name: &'static str,
364-
module: Option<&'static str>,
364+
module: &'static str,
365365
},
366366
Map(HashMap<&'static str, IntrospectionNode<'a>>),
367367
List(Vec<AttributedIntrospectionNode<'a>>),
@@ -425,11 +425,8 @@ impl IntrospectionNode<'_> {
425425
content.push_tokens(serialize_type_hint(annotation, pyo3_crate_path));
426426
}
427427
Self::ConstantType { name, module } => {
428-
let annotation = if let Some(module) = module {
429-
quote! { #pyo3_crate_path::inspect::TypeHint::module_attr(#module, #name) }
430-
} else {
431-
quote! { #pyo3_crate_path::inspect::TypeHint::builtin(#name) }
432-
};
428+
let annotation =
429+
quote! { #pyo3_crate_path::inspect::TypeHint::module_attr(#module, #name) };
433430
content.push_tokens(serialize_type_hint(annotation, pyo3_crate_path));
434431
}
435432
Self::Map(map) => {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -443,7 +443,7 @@ fn get_class_type_hint(cls: &Ident, args: &PyClassArgs, ctx: &Ctx) -> TokenStrea
443443
let module = module.value.value();
444444
quote! { #pyo3_path::inspect::TypeHint::module_attr(#module, #name) }
445445
} else {
446-
quote! { #pyo3_path::inspect::TypeHint::builtin(#name) }
446+
quote! { #pyo3_path::inspect::TypeHint::local(#name) }
447447
}
448448
}
449449

‎src/inspect/mod.rs‎

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ pub struct TypeHint {
2424

2525
#[derive(Clone, Copy)]
2626
enum TypeHintExpr {
27+
/// A local name. Used only when the module is unknown.
28+
Local { id: &'static str },
2729
/// A built-name like `list` or `datetime`. Used for built-in types or modules.
2830
Builtin { id: &'static str },
2931
/// A module member like `datetime.time` where module = `datetime` and attr = `time`
@@ -65,7 +67,19 @@ impl TypeHint {
6567
/// ```
6668
pub const fn module_attr(module: &'static str, attr: &'static str) -> Self {
6769
Self {
68-
inner: TypeHintExpr::ModuleAttribute { module, attr },
70+
inner: if matches!(module.as_bytes(), b"builtins") {
71+
TypeHintExpr::Builtin { id: attr }
72+
} else {
73+
TypeHintExpr::ModuleAttribute { module, attr }
74+
},
75+
}
76+
}
77+
78+
/// A value in the local module which module is unknown
79+
#[doc(hidden)]
80+
pub const fn local(name: &'static str) -> Self {
81+
Self {
82+
inner: TypeHintExpr::Local { id: name },
6983
}
7084
}
7185

@@ -105,6 +119,11 @@ impl TypeHint {
105119
pub const fn serialize_for_introspection(hint: &TypeHint, mut output: &mut [u8]) -> usize {
106120
let original_len = output.len();
107121
match &hint.inner {
122+
TypeHintExpr::Local { id } => {
123+
output = write_slice_and_move_forward(b"{\"type\":\"local\",\"id\":\"", output);
124+
output = write_slice_and_move_forward(id.as_bytes(), output);
125+
output = write_slice_and_move_forward(b"\"}", output);
126+
}
108127
TypeHintExpr::Builtin { id } => {
109128
output = write_slice_and_move_forward(b"{\"type\":\"builtin\",\"id\":\"", output);
110129
output = write_slice_and_move_forward(id.as_bytes(), output);
@@ -151,6 +170,7 @@ pub const fn serialize_for_introspection(hint: &TypeHint, mut output: &mut [u8])
151170
#[doc(hidden)]
152171
pub const fn serialized_len_for_introspection(hint: &TypeHint) -> usize {
153172
match &hint.inner {
173+
TypeHintExpr::Local { id } => 24 + id.len(),
154174
TypeHintExpr::Builtin { id } => 26 + id.len(),
155175
TypeHintExpr::ModuleAttribute { module, attr } => 42 + module.len() + attr.len(),
156176
TypeHintExpr::Union { elts } => {
@@ -184,7 +204,7 @@ impl fmt::Display for TypeHint {
184204
#[inline]
185205
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
186206
match &self.inner {
187-
TypeHintExpr::Builtin { id } => id.fmt(f),
207+
TypeHintExpr::Builtin { id } | TypeHintExpr::Local { id } => id.fmt(f),
188208
TypeHintExpr::ModuleAttribute { module, attr } => {
189209
module.fmt(f)?;
190210
f.write_str(".")?;
@@ -241,19 +261,23 @@ mod tests {
241261
const T: TypeHint = TypeHint::subscript(
242262
&TypeHint::builtin("dict"),
243263
&[
244-
TypeHint::union(&[TypeHint::builtin("int"), TypeHint::builtin("float")]),
264+
TypeHint::union(&[
265+
TypeHint::builtin("int"),
266+
TypeHint::module_attr("builtins", "float"),
267+
TypeHint::local("weird"),
268+
]),
245269
TypeHint::module_attr("datetime", "time"),
246270
],
247271
);
248-
assert_eq!(T.to_string(), "dict[int | float, datetime.time]")
272+
assert_eq!(T.to_string(), "dict[int | float | weird, datetime.time]")
249273
}
250274

251275
#[test]
252276
fn test_serialize_for_introspection() {
253277
const T: TypeHint = TypeHint::subscript(
254278
&TypeHint::builtin("dict"),
255279
&[
256-
TypeHint::union(&[TypeHint::builtin("int"), TypeHint::builtin("float")]),
280+
TypeHint::union(&[TypeHint::builtin("int"), TypeHint::local("weird")]),
257281
TypeHint::module_attr("datetime", "time"),
258282
],
259283
);
@@ -265,7 +289,7 @@ mod tests {
265289
};
266290
assert_eq!(
267291
std::str::from_utf8(&SER).unwrap(),
268-
r#"{"type":"subscript","value":{"type":"builtin","id":"dict"},"slice":[{"type":"union","elts":[{"type":"builtin","id":"int"},{"type":"builtin","id":"float"}]},{"type":"attribute","module":"datetime","attr":"time"}]}"#
292+
r#"{"type":"subscript","value":{"type":"builtin","id":"dict"},"slice":[{"type":"union","elts":[{"type":"builtin","id":"int"},{"type":"local","id":"weird"}]},{"type":"attribute","module":"datetime","attr":"time"}]}"#
269293
)
270294
}
271295
}

0 commit comments

Comments
 (0)