From a8d9104fa3a1351485fd0e41b0f28501e6890dd4 Mon Sep 17 00:00:00 2001 From: Simon Date: Thu, 19 Sep 2024 10:13:37 +0200 Subject: [PATCH] Fix/#13070 defer annotations when future is active (#13395) --- .../src/semantic_index.rs | 9 +++ .../src/semantic_index/builder.rs | 15 ++++ .../src/types/infer.rs | 76 ++++++++++++++++--- 3 files changed, 90 insertions(+), 10 deletions(-) diff --git a/crates/red_knot_python_semantic/src/semantic_index.rs b/crates/red_knot_python_semantic/src/semantic_index.rs index 8019698198928..5ec6d57ccb112 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..56df1c44d9ade 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,16 @@ where &alias.name.id }; + // 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 == "annotations" + && node.module.as_deref() == Some("__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 c5a502307fc31..fa5197f5f3e9b 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.index.has_future_annotations() || self.file.is_stub(self.db.upcast()) } /// Are we currently inferring deferred types? @@ -703,7 +704,7 @@ impl<'db> TypeInferenceBuilder<'db> { self.infer_parameters(parameters); // TODO: this should also be applied to parameter annotations. - if self.is_stub() { + if self.are_all_types_deferred() { self.types.has_deferred = true; } else { self.infer_optional_annotation_expression(returns.as_deref()); @@ -831,9 +832,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.are_all_types_deferred() { self.types.has_deferred = true; } else { for base in class.bases() { @@ -843,13 +844,11 @@ 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()); } fn infer_class_deferred(&mut self, class: &ast::StmtClassDef) { - if self.is_stub() { + if self.are_all_types_deferred() { for base in class.bases() { self.infer_expression(base); } @@ -4166,6 +4165,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();