From 3ceedf76b8df1cfe749d0bb4c4f0aad6ee97ef6a Mon Sep 17 00:00:00 2001
From: Dylan <53534755+dylwil3@users.noreply.github.com>
Date: Fri, 30 Aug 2024 18:48:06 -0500
Subject: [PATCH] [red-knot] Infer type of class constructor call expression
 (#13171)

This tiny PR implements the following type inference: the type of
`Foo(...)` will be `Foo`.

---------

Co-authored-by: Carl Meyer <carl@astral.sh>
---
 crates/red_knot_python_semantic/src/types.rs   |  4 ++--
 .../src/types/infer.rs                         | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
index c1a270c54e3af..d7b9c8e91d4d3 100644
--- a/crates/red_knot_python_semantic/src/types.rs
+++ b/crates/red_knot_python_semantic/src/types.rs
@@ -328,8 +328,8 @@ impl<'db> Type<'db> {
         match self {
             Type::Function(function_type) => function_type.returns(db).or(Some(Type::Unknown)),
 
-            // TODO: handle class constructors
-            Type::Class(_class_ty) => Some(Type::Unknown),
+            // TODO annotated return type on `__new__` or metaclass `__call__`
+            Type::Class(class) => Some(Type::Instance(*class)),
 
             // TODO: handle classes which implement the Callable protocol
             Type::Instance(_instance_ty) => Some(Type::Unknown),
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index f8d73108ae354..5286fa0420073 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -2834,6 +2834,24 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn class_constructor_call_expression() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_dedented(
+            "src/a.py",
+            "
+            class Foo: ...
+
+            x = Foo()
+            ",
+        )?;
+
+        assert_public_ty(&db, "src/a.py", "x", "Foo");
+
+        Ok(())
+    }
+
     #[test]
     fn resolve_union() -> anyhow::Result<()> {
         let mut db = setup_db();