[red-knot] fix scope inference with deferred types (#13204)

Test coverage for #13131 wasn't as good as I thought it was, because
although we infer a lot of types in stubs in typeshed, we don't check
typeshed, and therefore we don't do scope-level inference and pull all
types for a scope. So we didn't really have good test coverage for
scope-level inference in a stub. And because of this, I got the code for
supporting that wrong, meaning that if we did scope-level inference with
deferred types, we'd end up never populating the deferred types in the
scope's `TypeInference`, which causes panics like #13160.

Here I both add test coverage by running the corpus tests both as `.py`
and as `.pyi` (which reveals the panic), and I fix the code to support
deferred types in scope inference.

This also revealed a problem with deferred types in generic functions,
which effectively span two scopes. That problem will require a bit more
thought, and I don't want to block this PR on it, so for now I just
don't defer annotations on generic functions.

Fixes #13160.

Cargo.lock

@@ -1997,6 +1997,7 @@ dependencies = [
  "ruff_text_size",
  "rustc-hash 2.0.0",
  "salsa",
+ "tempfile",
  "tracing",
 ]
 

crates/red_knot_python_semantic/src/types/infer.rs

@@ -468,11 +468,8 @@ impl<'db> TypeInferenceBuilder<'db> {
             .as_deref()
             .expect("function type params scope without type params");
 
-        // TODO: this should also be applied to parameter annotations.
-        if !self.is_stub() {
-            self.infer_optional_expression(function.returns.as_deref());
-        }
-
+        // TODO: defer annotation resolution in stubs, with __future__.annotations, or stringified
+        self.infer_optional_expression(function.returns.as_deref());
         self.infer_type_parameters(type_params);
         self.infer_parameters(&function.parameters);
     }
@@ -567,7 +564,9 @@ 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.types.has_deferred = true;
+            } else {
                 self.infer_optional_annotation_expression(returns.as_deref());
             }
         }
@@ -684,7 +683,9 @@ impl<'db> TypeInferenceBuilder<'db> {
 
         // inference of bases deferred in stubs
         // TODO also defer stringified generic type parameters
-        if !self.is_stub() {
+        if self.is_stub() {
+            self.types.has_deferred = true;
+        } else {
             for base in class.bases() {
                 self.infer_expression(base);
             }
@@ -693,14 +694,12 @@ impl<'db> TypeInferenceBuilder<'db> {
 
     fn infer_function_deferred(&mut self, function: &ast::StmtFunctionDef) {
         if self.is_stub() {
-            self.types.has_deferred = true;
             self.infer_optional_annotation_expression(function.returns.as_deref());
         }
     }
 
     fn infer_class_deferred(&mut self, class: &ast::StmtClassDef) {
         if self.is_stub() {
-            self.types.has_deferred = true;
             for base in class.bases() {
                 self.infer_expression(base);
             }

crates/red_knot_workspace/Cargo.toml

@@ -29,6 +29,7 @@ tracing = { workspace = true }
 
 [dev-dependencies]
 ruff_db = { workspace = true, features = ["testing"] }
+tempfile = { workspace = true }
 
 [lints]
 workspace = true

crates/red_knot_workspace/tests/check.rs

@@ -21,19 +21,27 @@ fn setup_db(workspace_root: &SystemPath) -> anyhow::Result<RootDatabase> {
 #[test]
 #[allow(clippy::print_stdout)]
 fn corpus_no_panic() -> anyhow::Result<()> {
-    let corpus = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("resources/test/corpus");
-    let system_corpus =
-        SystemPathBuf::from_path_buf(corpus.clone()).expect("corpus path to be UTF8");
-    let db = setup_db(&system_corpus)?;
+    let root = SystemPathBuf::from_path_buf(tempfile::TempDir::new()?.into_path()).unwrap();
+    let db = setup_db(&root)?;
 
-    for path in fs::read_dir(&corpus).expect("corpus to be a directory") {
-        let path = path.expect("path to not be an error").path();
-        println!("checking {path:?}");
-        let path = SystemPathBuf::from_path_buf(path.clone()).expect("path to be UTF-8");
+    let corpus = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("resources/test/corpus");
+
+    for path in fs::read_dir(&corpus)? {
+        let source = path?.path();
+        println!("checking {source:?}");
+        let source_fn = source.file_name().unwrap().to_str().unwrap();
+        let py_dest = root.join(source_fn);
+        fs::copy(&source, py_dest.as_std_path())?;
         // this test is only asserting that we can pull every expression type without a panic
         // (and some non-expressions that clearly define a single type)
-        let file = system_path_to_file(&db, path).expect("file to exist");
+        let file = system_path_to_file(&db, py_dest).unwrap();
+        pull_types(&db, file);
 
+        // try the file as a stub also
+        println!("re-checking as .pyi");
+        let pyi_dest = root.join(format!("{source_fn}i"));
+        std::fs::copy(source, pyi_dest.as_std_path())?;
+        let file = system_path_to_file(&db, pyi_dest).unwrap();
         pull_types(&db, file);
     }
     Ok(())