From 7cefe1855f3efb8cdcabd486f28c1ea814fc21fb Mon Sep 17 00:00:00 2001 From: slyces Date: Wed, 18 Sep 2024 16:58:10 +0200 Subject: [PATCH 1/4] [red-knot] feat: resolve deferred type hints with __future__ annotations Fixes #13070 --- .../src/semantic_index.rs | 9 +++ .../src/semantic_index/builder.rs | 14 ++++ .../src/types/infer.rs | 74 +++++++++++++++++-- 3 files changed, 90 insertions(+), 7 deletions(-) diff --git a/crates/red_knot_python_semantic/src/semantic_index.rs b/crates/red_knot_python_semantic/src/semantic_index.rs index 1d1700c765fba..a0d0d42a702b0 100644 --- a/crates/red_knot_python_semantic/src/semantic_index.rs +++ b/crates/red_knot_python_semantic/src/semantic_index.rs @@ -115,6 +115,9 @@ pub(crate) struct SemanticIndex<'db> { /// Note: We should not depend on this map when analysing other files or /// changing a file invalidates all dependents. ast_ids: IndexVec, + + /// Flags about the global scope (code usage impacting inference) + has_future_annotations: bool, } impl<'db> SemanticIndex<'db> { @@ -215,6 +218,12 @@ impl<'db> SemanticIndex<'db> { pub(crate) fn node_scope(&self, node: NodeWithScopeRef) -> FileScopeId { self.scopes_by_node[&node.node_key()] } + + /// Checks if there is an import of `__future__.annotations` in the global scope, which affects + /// the logic for type inference. + pub(super) fn has_future_annotations(&self) -> bool { + self.has_future_annotations + } } pub struct AncestorsIter<'a> { diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index d73f554bd047a..bd87af5b4b6a9 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -45,6 +45,9 @@ pub(super) struct SemanticIndexBuilder<'db> { /// Flow states at each `break` in the current loop. loop_break_states: Vec, + /// Flags about the file's global scope + has_future_annotations: bool, + // Semantic Index fields scopes: IndexVec, scope_ids_by_scope: IndexVec>, @@ -68,6 +71,8 @@ impl<'db> SemanticIndexBuilder<'db> { current_match_case: None, loop_break_states: vec![], + has_future_annotations: false, + scopes: IndexVec::new(), symbol_tables: IndexVec::new(), ast_ids: IndexVec::new(), @@ -450,6 +455,7 @@ impl<'db> SemanticIndexBuilder<'db> { scopes_by_expression: self.scopes_by_expression, scopes_by_node: self.scopes_by_node, use_def_maps, + has_future_annotations: self.has_future_annotations, } } } @@ -543,7 +549,15 @@ where &alias.name.id }; + // Look for imports `from __future__ import annotations`, ignore `as ...` + self.has_future_annotations |= alias.name.id.as_str() == "annotations" + && node + .module + .as_ref() + .is_some_and(|module| module.id().as_str() == "__future__"); + let symbol = self.add_symbol(symbol_name.clone()); + self.add_definition(symbol, ImportFromDefinitionNodeRef { node, alias_index }); } } diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index a38e6cc6bd6b2..0adf97fb9d921 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -328,6 +328,10 @@ impl<'db> TypeInferenceBuilder<'db> { matches!(self.region, InferenceRegion::Deferred(_)) } + fn has_future_annotations(&self) -> bool { + self.index.has_future_annotations() + } + /// Infers types in the given [`InferenceRegion`]. fn infer_region(&mut self) { match self.region { @@ -697,7 +701,7 @@ impl<'db> TypeInferenceBuilder<'db> { self.infer_parameters(parameters); // TODO: this should also be applied to parameter annotations. - if self.is_stub() { + if self.is_stub() || self.has_future_annotations() { self.types.has_deferred = true; } else { self.infer_optional_annotation_expression(returns.as_deref()); @@ -825,9 +829,9 @@ impl<'db> TypeInferenceBuilder<'db> { self.infer_expression(&keyword.value); } - // inference of bases deferred in stubs + // Inference of bases deferred in stubs // TODO also defer stringified generic type parameters - if self.is_stub() { + if self.is_stub() || self.has_future_annotations() { self.types.has_deferred = true; } else { for base in class.bases() { @@ -837,13 +841,12 @@ impl<'db> TypeInferenceBuilder<'db> { } fn infer_function_deferred(&mut self, function: &ast::StmtFunctionDef) { - if self.is_stub() { - self.infer_optional_annotation_expression(function.returns.as_deref()); - } + self.infer_optional_annotation_expression(function.returns.as_deref()); + if self.is_stub() || self.has_future_annotations() {} } fn infer_class_deferred(&mut self, class: &ast::StmtClassDef) { - if self.is_stub() { + if self.is_stub() || self.has_future_annotations() { for base in class.bases() { self.infer_expression(base); } @@ -4009,6 +4012,63 @@ mod tests { Ok(()) } + #[test] + fn deferred_annotation_in_stubs_always_resolve() -> anyhow::Result<()> { + let mut db = setup_db(); + + // Stub files should always resolve deferred annotations + db.write_dedented( + "/src/stub.pyi", + " + def get_foo() -> Foo: ... + class Foo: ... + foo = get_foo() + ", + )?; + assert_public_ty(&db, "/src/stub.pyi", "foo", "Foo"); + + Ok(()) + } + + #[test] + fn deferred_annotations_regular_source_fails() -> anyhow::Result<()> { + let mut db = setup_db(); + + // In (regular) source files, deferred annotations are *not* resolved + // Also tests imports from `__future__` that are not annotations + db.write_dedented( + "/src/source.py", + " + from __future__ import with_statement as annotations + def get_foo() -> Foo: ... + class Foo: ... + foo = get_foo() + ", + )?; + assert_public_ty(&db, "/src/source.py", "foo", "Unknown"); + + Ok(()) + } + + #[test] + fn deferred_annotation_in_sources_with_future_resolves() -> anyhow::Result<()> { + let mut db = setup_db(); + + // In source files with `__future__.annotations`, deferred annotations are resolved + db.write_dedented( + "/src/source_with_future.py", + " + from __future__ import annotations + def get_foo() -> Foo: ... + class Foo: ... + foo = get_foo() + ", + )?; + assert_public_ty(&db, "/src/source_with_future.py", "foo", "Foo"); + + Ok(()) + } + #[test] fn narrow_not_none() -> anyhow::Result<()> { let mut db = setup_db(); From 75bed691d760b5f11576804885c93b6b95abbdb7 Mon Sep 17 00:00:00 2001 From: slyces Date: Wed, 18 Sep 2024 18:04:07 +0200 Subject: [PATCH 2/4] [red-knot] refactor: condense usages of `is_stubs` & `has_future_annotations` --- .../src/types/infer.rs | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 0adf97fb9d921..6acf0c3ba73e2 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -318,9 +318,10 @@ impl<'db> TypeInferenceBuilder<'db> { self.types.has_deferred |= inference.has_deferred; } - /// Are we currently inferring types in a stub file? - fn is_stub(&self) -> bool { - self.file.is_stub(self.db.upcast()) + /// Are we currently inferring types in file with deferred types? + /// This is true for stub files and files with `__future__.annotations` + fn are_all_types_deferred(&self) -> bool { + self.file.is_stub(self.db.upcast()) || self.index.has_future_annotations() } /// Are we currently inferring deferred types? @@ -328,10 +329,6 @@ impl<'db> TypeInferenceBuilder<'db> { matches!(self.region, InferenceRegion::Deferred(_)) } - fn has_future_annotations(&self) -> bool { - self.index.has_future_annotations() - } - /// Infers types in the given [`InferenceRegion`]. fn infer_region(&mut self) { match self.region { @@ -701,7 +698,7 @@ impl<'db> TypeInferenceBuilder<'db> { self.infer_parameters(parameters); // TODO: this should also be applied to parameter annotations. - if self.is_stub() || self.has_future_annotations() { + if self.are_all_types_deferred() { self.types.has_deferred = true; } else { self.infer_optional_annotation_expression(returns.as_deref()); @@ -831,7 +828,7 @@ impl<'db> TypeInferenceBuilder<'db> { // Inference of bases deferred in stubs // TODO also defer stringified generic type parameters - if self.is_stub() || self.has_future_annotations() { + if self.are_all_types_deferred() { self.types.has_deferred = true; } else { for base in class.bases() { @@ -842,11 +839,10 @@ impl<'db> TypeInferenceBuilder<'db> { fn infer_function_deferred(&mut self, function: &ast::StmtFunctionDef) { self.infer_optional_annotation_expression(function.returns.as_deref()); - if self.is_stub() || self.has_future_annotations() {} } fn infer_class_deferred(&mut self, class: &ast::StmtClassDef) { - if self.is_stub() || self.has_future_annotations() { + if self.are_all_types_deferred() { for base in class.bases() { self.infer_expression(base); } From e4e033167bb6f11b8b48f3bb073ecdda37bf718c Mon Sep 17 00:00:00 2001 From: slyces Date: Wed, 18 Sep 2024 19:29:14 +0200 Subject: [PATCH 3/4] fixup! [red-knot] feat: resolve deferred type hints with __future__ annotations --- crates/red_knot_python_semantic/src/semantic_index/builder.rs | 4 ++++ crates/red_knot_python_semantic/src/types/infer.rs | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index bd87af5b4b6a9..8d99f8ceb3655 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -550,6 +550,10 @@ where }; // Look for imports `from __future__ import annotations`, ignore `as ...` + // We intentionally don't enforce the rules about location of `__future__` + // imports here, we assume the user's intent was to apply the `__future__` + // import, so we still check using it (and will also emit a diagnostic about a + // miss-placed `__future__` import.) self.has_future_annotations |= alias.name.id.as_str() == "annotations" && node .module diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 6acf0c3ba73e2..751f30a991532 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -321,7 +321,7 @@ impl<'db> TypeInferenceBuilder<'db> { /// Are we currently inferring types in file with deferred types? /// This is true for stub files and files with `__future__.annotations` fn are_all_types_deferred(&self) -> bool { - self.file.is_stub(self.db.upcast()) || self.index.has_future_annotations() + self.index.has_future_annotations() || self.file.is_stub(self.db.upcast()) } /// Are we currently inferring deferred types? From 301d0651284084e701fdea4787d4cebf372e38a4 Mon Sep 17 00:00:00 2001 From: slyces Date: Thu, 19 Sep 2024 10:02:03 +0200 Subject: [PATCH 4/4] fixup! [red-knot] feat: resolve deferred type hints with __future__ annotations --- .../red_knot_python_semantic/src/semantic_index/builder.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index 8d99f8ceb3655..56df1c44d9ade 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -554,11 +554,8 @@ where // imports here, we assume the user's intent was to apply the `__future__` // import, so we still check using it (and will also emit a diagnostic about a // miss-placed `__future__` import.) - self.has_future_annotations |= alias.name.id.as_str() == "annotations" - && node - .module - .as_ref() - .is_some_and(|module| module.id().as_str() == "__future__"); + self.has_future_annotations |= alias.name.id == "annotations" + && node.module.as_deref() == Some("__future__"); let symbol = self.add_symbol(symbol_name.clone());