Fix/#13070 defer annotations when future is active (#13395)

Author: Slyces

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<FileScopeId, AstIds>,
+
+    /// 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> {

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<FlowSnapshot>,
 
+    /// Flags about the file's global scope
+    has_future_annotations: bool,
+
     // Semantic Index fields
     scopes: IndexVec<FileScopeId, Scope>,
     scope_ids_by_scope: IndexVec<FileScopeId, ScopeId<'db>>,
@@ -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 });
                 }
             }

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();