Address review feedback from 11963 (#12145)

crates/red_knot_python_semantic/src/semantic_model.rs

@@ -5,7 +5,7 @@ use ruff_python_ast::{Expr, ExpressionRef, StmtClassDef};
 
 use crate::semantic_index::ast_ids::HasScopedAstId;
 use crate::semantic_index::definition::Definition;
-use crate::semantic_index::symbol::{PublicSymbolId, ScopeKind};
+use crate::semantic_index::symbol::PublicSymbolId;
 use crate::semantic_index::{public_symbol, semantic_index, NodeWithScopeKey};
 use crate::types::{infer_types, public_symbol_ty, Type, TypingContext};
 use crate::Db;
@@ -38,6 +38,10 @@ impl<'db> SemanticModel<'db> {
 }
 
 pub trait HasTy {
+    /// Returns the inferred type of `self`.
+    ///
+    /// ## Panics
+    /// May panic if `self` is from another file than `model`.
     fn ty<'db>(&self, model: &SemanticModel<'db>) -> Type<'db>;
 }
 
@@ -48,7 +52,6 @@ impl HasTy for ast::ExpressionRef<'_> {
         let scope = file_scope.to_scope_id(model.db, model.file);
 
         let expression_id = self.scoped_ast_id(model.db, scope);
-
         infer_types(model.db, scope).expression_ty(expression_id)
     }
 }
@@ -142,15 +145,7 @@ impl HasTy for ast::StmtFunctionDef {
         let index = semantic_index(model.db, model.file);
         let definition_scope = index.definition_scope(NodeWithScopeKey::from(self));
 
-        // SAFETY: A function always has either an enclosing module, function or class scope.
-        let mut parent_scope_id = index.parent_scope_id(definition_scope).unwrap();
-        let parent_scope = index.scope(parent_scope_id);
-
-        if parent_scope.kind() == ScopeKind::Annotation {
-            parent_scope_id = index.parent_scope_id(parent_scope_id).unwrap();
-        }
-
-        let scope = parent_scope_id.to_scope_id(model.db, model.file);
+        let scope = definition_scope.to_scope_id(model.db, model.file);
 
         let types = infer_types(model.db, scope);
         let definition = Definition::FunctionDef(self.scoped_ast_id(model.db, scope));
@@ -163,16 +158,7 @@ impl HasTy for StmtClassDef {
     fn ty<'db>(&self, model: &SemanticModel<'db>) -> Type<'db> {
         let index = semantic_index(model.db, model.file);
         let definition_scope = index.definition_scope(NodeWithScopeKey::from(self));
-
-        // SAFETY: A class always has either an enclosing module, function or class scope.
-        let mut parent_scope_id = index.parent_scope_id(definition_scope).unwrap();
-        let parent_scope = index.scope(parent_scope_id);
-
-        if parent_scope.kind() == ScopeKind::Annotation {
-            parent_scope_id = index.parent_scope_id(parent_scope_id).unwrap();
-        }
-
-        let scope = parent_scope_id.to_scope_id(model.db, model.file);
+        let scope = definition_scope.to_scope_id(model.db, model.file);
 
         let types = infer_types(model.db, scope);
         let definition = Definition::ClassDef(self.scoped_ast_id(model.db, scope));
@@ -180,3 +166,68 @@ impl HasTy for StmtClassDef {
         types.definition_ty(definition)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use red_knot_module_resolver::{set_module_resolution_settings, ModuleResolutionSettings};
+    use ruff_db::file_system::FileSystemPathBuf;
+    use ruff_db::parsed::parsed_module;
+    use ruff_db::vfs::system_path_to_file;
+
+    use crate::db::tests::TestDb;
+    use crate::types::Type;
+    use crate::{HasTy, SemanticModel};
+
+    fn setup_db() -> TestDb {
+        let mut db = TestDb::new();
+        set_module_resolution_settings(
+            &mut db,
+            ModuleResolutionSettings {
+                extra_paths: vec![],
+                workspace_root: FileSystemPathBuf::from("/src"),
+                site_packages: None,
+                custom_typeshed: None,
+            },
+        );
+
+        db
+    }
+
+    #[test]
+    fn function_ty() -> anyhow::Result<()> {
+        let db = setup_db();
+
+        db.memory_file_system()
+            .write_file("/src/foo.py", "def test(): pass")?;
+        let foo = system_path_to_file(&db, "/src/foo.py").unwrap();
+
+        let ast = parsed_module(&db, foo);
+
+        let function = ast.suite()[0].as_function_def_stmt().unwrap();
+        let model = SemanticModel::new(&db, foo);
+        let ty = function.ty(&model);
+
+        assert!(matches!(ty, Type::Function(_)));
+
+        Ok(())
+    }
+
+    #[test]
+    fn class_ty() -> anyhow::Result<()> {
+        let db = setup_db();
+
+        db.memory_file_system()
+            .write_file("/src/foo.py", "class Test: pass")?;
+        let foo = system_path_to_file(&db, "/src/foo.py").unwrap();
+
+        let ast = parsed_module(&db, foo);
+
+        let class = ast.suite()[0].as_class_def_stmt().unwrap();
+        let model = SemanticModel::new(&db, foo);
+        let ty = class.ty(&model);
+
+        assert!(matches!(ty, Type::Class(_)));
+
+        Ok(())
+    }
+}