[red-knot] Handle context managers in (sync) with statements

crates/red_knot_python_semantic/resources/mdtest/declaration/async_with.md

@@ -0,0 +1,20 @@
+# Async with declarations
+
+## Basic `async with` statement
+
+The type of the target variable in a `with` statement should be the return type from the context manager's `__aenter__` method.
+However, Knot doesn't support `async with` statements yet. This test asserts that it doesn't emit any context manager-related errors.
+
+```py
+class Target: ...
+
+class Manager:
+    async def __aenter__(self) -> Target:
+        return Target()
+
+    async def __aexit__(self, exc_type, exc_value, traceback): ...
+
+async def test():
+    async with Manager() as f:
+        reveal_type(f)  # revealed: @Todo
+```

crates/red_knot_python_semantic/resources/mdtest/declaration/with.md

@@ -0,0 +1,140 @@
+# With declarations
+
+## Basic `with` statement
+
+The type of the target variable in a `with` statement is the return type from the context manager's `__enter__` method.
+
+```py
+class Target: ...
+
+class Manager:
+    def __enter__(self) -> Target:
+        return Target()
+
+    def __exit__(self, exc_type, exc_value, traceback): ...
+
+with Manager() as f:
+    reveal_type(f)  # revealed: Target
+```
+
+## Union context manager
+
+```py
+def coinflip() -> bool:
+    return True
+
+class Manager1:
+    def __enter__(self) -> str:
+        return "foo"
+
+    def __exit__(self, exc_type, exc_value, traceback): ...
+
+class Manager2:
+    def __enter__(self) -> int:
+        return 42
+
+    def __exit__(self, exc_type, exc_value, traceback): ...
+
+context_expr = Manager1() if coinflip() else Manager2()
+
+with context_expr as f:
+    reveal_type(f)  # revealed: str | int
+```
+
+## Context manager without an `__enter__` or `__exit__` method
+
+```py
+class Manager: ...
+
+# error: [invalid-context-manager] "Object of type Manager cannot be used with `with` because it doesn't implement `__enter__` and `__exit__`"
+with Manager():
+    ...
+```
+
+## Context manager without an `__enter__` method
+
+```py
+class Manager:
+    def __exit__(self, exc_tpe, exc_value, traceback): ...
+
+# error: [invalid-context-manager] "Object of type Manager cannot be used with `with` because it doesn't implement `__enter__`"
+with Manager():
+    ...
+```
+
+## Context manager without an `__exit__` method
+
+```py
+class Manager:
+    def __enter__(self): ...
+
+# error: [invalid-context-manager] "Object of type Manager cannot be used with `with` because it doesn't implement `__exit__`"
+with Manager():
+    ...
+```
+
+## Context manager with non-callable `__enter__` attribute
+
+```py
+class Manager:
+    __enter__ = 42
+
+    def __exit__(self, exc_tpe, exc_value, traceback): ...
+
+# error: [invalid-context-manager] "Object of type Manager cannot be used with `with` because the method `__enter__` of type Literal[42] is not callable"
+with Manager():
+    ...
+```
+
+## Context manager with non-callable `__exit__` attribute
+
+```py
+class Manager:
+    def __enter__(self) -> Self: ...
+
+    __exit__ = 32
+
+# error: [invalid-context-manager] "Object of type Manager cannot be used with `with` because the method `__exit__` of type Literal[32] is not callable"
+with Manager():
+    ...
+```
+
+## Context expression with non-callable union variants
+
+```py
+def coinflip() -> bool:
+    return True
+
+class Manager1:
+    def __enter__(self) -> str:
+        return "foo"
+
+    def __exit__(self, exc_type, exc_value, traceback): ...
+
+class NotAContextManager: ...
+
+context_expr = Manager1() if coinflip() else NotAContextManager()
+
+# error: [invalid-context-manager] "Object of type Manager1 | NotAContextManager cannot be used with `with` because the method `__enter__` of type Literal[__enter__] | Unbound is not callable"
+# error: [invalid-context-manager] "Object of type Manager1 | NotAContextManager cannot be used with `with` because the method `__exit__` of type Literal[__exit__] | Unbound is not callable"
+with context_expr as f:
+    reveal_type(f)  # revealed: str | Unknown
+```
+
+## Context expression with "sometimes" callable `__enter__` method
+
+```py
+def coinflip() -> bool:
+    return True
+
+class Manager:
+    if coinflip():
+        def __enter__(self) -> str:
+            return "abcd"
+
+    def __exit__(self, *args): ...
+
+with Manager() as f:
+    # TODO: This should emit an error that `__enter__` is possibly unbound.
+    reveal_type(f)  # revealed: str
+```

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -734,12 +734,20 @@ where
                     self.flow_merge(break_state);
                 }
             }
-            ast::Stmt::With(ast::StmtWith { items, body, .. }) => {
+            ast::Stmt::With(ast::StmtWith {
+                items,
+                body,
+                is_async,
+                ..
+            }) => {
                 for item in items {
                     self.visit_expr(&item.context_expr);
                     if let Some(optional_vars) = item.optional_vars.as_deref() {
                         self.add_standalone_expression(&item.context_expr);
-                        self.push_assignment(item.into());
+                        self.push_assignment(CurrentAssignment::WithItem {
+                            item,
+                            is_async: *is_async,
+                        });
                         self.visit_expr(optional_vars);
                         self.pop_assignment();
                     }
@@ -1011,12 +1019,13 @@ where
                                 },
                             );
                         }
-                        Some(CurrentAssignment::WithItem(with_item)) => {
+                        Some(CurrentAssignment::WithItem { item, is_async }) => {
                             self.add_definition(
                                 symbol,
                                 WithItemDefinitionNodeRef {
-                                    node: with_item,
+                                    node: item,
                                     target: name_node,
+                                    is_async,
                                 },
                             );
                         }
@@ -1232,7 +1241,10 @@ enum CurrentAssignment<'a> {
         node: &'a ast::Comprehension,
         first: bool,
     },
-    WithItem(&'a ast::WithItem),
+    WithItem {
+        item: &'a ast::WithItem,
+        is_async: bool,
+    },
 }
 
 impl<'a> From<&'a ast::StmtAnnAssign> for CurrentAssignment<'a> {
@@ -1259,12 +1271,6 @@ impl<'a> From<&'a ast::ExprNamed> for CurrentAssignment<'a> {
     }
 }
 
-impl<'a> From<&'a ast::WithItem> for CurrentAssignment<'a> {
-    fn from(value: &'a ast::WithItem) -> Self {
-        Self::WithItem(value)
-    }
-}
-
 struct CurrentMatchCase<'a> {
     /// The pattern that's part of the current match case.
     pattern: &'a ast::Pattern,

crates/red_knot_python_semantic/src/semantic_index/definition.rs

@@ -176,6 +176,7 @@ pub(crate) struct AssignmentDefinitionNodeRef<'a> {
 pub(crate) struct WithItemDefinitionNodeRef<'a> {
     pub(crate) node: &'a ast::WithItem,
     pub(crate) target: &'a ast::ExprName,
+    pub(crate) is_async: bool,
 }
 
 #[derive(Copy, Clone, Debug)]
@@ -277,12 +278,15 @@ impl DefinitionNodeRef<'_> {
                     DefinitionKind::ParameterWithDefault(AstNodeRef::new(parsed, parameter))
                 }
             },
-            DefinitionNodeRef::WithItem(WithItemDefinitionNodeRef { node, target }) => {
-                DefinitionKind::WithItem(WithItemDefinitionKind {
-                    node: AstNodeRef::new(parsed.clone(), node),
-                    target: AstNodeRef::new(parsed, target),
-                })
-            }
+            DefinitionNodeRef::WithItem(WithItemDefinitionNodeRef {
+                node,
+                target,
+                is_async,
+            }) => DefinitionKind::WithItem(WithItemDefinitionKind {
+                node: AstNodeRef::new(parsed.clone(), node),
+                target: AstNodeRef::new(parsed, target),
+                is_async,
+            }),
             DefinitionNodeRef::MatchPattern(MatchPatternDefinitionNodeRef {
                 pattern,
                 identifier,
@@ -329,7 +333,11 @@ impl DefinitionNodeRef<'_> {
                 ast::AnyParameterRef::Variadic(parameter) => parameter.into(),
                 ast::AnyParameterRef::NonVariadic(parameter) => parameter.into(),
             },
-            Self::WithItem(WithItemDefinitionNodeRef { node: _, target }) => target.into(),
+            Self::WithItem(WithItemDefinitionNodeRef {
+                node: _,
+                target,
+                is_async: _,
+            }) => target.into(),
             Self::MatchPattern(MatchPatternDefinitionNodeRef { identifier, .. }) => {
                 identifier.into()
             }
@@ -534,6 +542,7 @@ pub enum AssignmentKind {
 pub struct WithItemDefinitionKind {
     node: AstNodeRef<ast::WithItem>,
     target: AstNodeRef<ast::ExprName>,
+    is_async: bool,
 }
 
 impl WithItemDefinitionKind {
@@ -544,6 +553,10 @@ impl WithItemDefinitionKind {
     pub(crate) fn target(&self) -> &ast::ExprName {
         self.target.node()
     }
+
+    pub(crate) const fn is_async(&self) -> bool {
+        self.is_async
+    }
 }
 
 #[derive(Clone, Debug)]

crates/red_knot_python_semantic/src/types.rs

@@ -1010,8 +1010,6 @@ impl<'db> Type<'db> {
     }
 
     /// Return the type resulting from calling an object of this type.
-    ///
-    /// Returns `None` if `self` is not a callable type.
     #[must_use]
     fn call(self, db: &'db dyn Db, arg_types: &[Type<'db>]) -> CallOutcome<'db> {
         match self {