From a1009b21c4fea1e9c8e9f77e86dc66566311bed7 Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Wed, 30 Oct 2024 14:38:17 +0100 Subject: [PATCH 1/3] [red-knot] Handle context managers in (sync) with statements --- .../resources/mdtest/declaration/with.md | 122 ++++++++++++++++++ crates/red_knot_python_semantic/src/types.rs | 2 - .../src/types/infer.rs | 112 ++++++++++++++-- 3 files changed, 225 insertions(+), 11 deletions(-) create mode 100644 crates/red_knot_python_semantic/resources/mdtest/declaration/with.md diff --git a/crates/red_knot_python_semantic/resources/mdtest/declaration/with.md b/crates/red_knot_python_semantic/resources/mdtest/declaration/with.md new file mode 100644 index 0000000000000..2b3bfa8f56eed --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/declaration/with.md @@ -0,0 +1,122 @@ +# 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 +``` diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs index 81c8ddba70f7c..9bb3242edc908 100644 --- a/crates/red_knot_python_semantic/src/types.rs +++ b/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 { diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index d62ee026471c3..9ab4cb65bd346 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -355,6 +355,7 @@ impl<'db> TypeInferenceBuilder<'db> { /// Get the already-inferred type of an expression node. /// /// PANIC if no type has been inferred for this node. + #[track_caller] fn expression_ty(&self, expr: &ast::Expr) -> Type<'db> { self.types .expression_ty(expr.scoped_ast_id(self.db, self.scope())) @@ -1003,13 +1004,20 @@ impl<'db> TypeInferenceBuilder<'db> { body, } = with_statement; + // TODO: Handle async with statements. For async with statements, the value returned + // by `__enter__` must be awaited. + for item in items { let target = item.optional_vars.as_deref(); if let Some(ast::Expr::Name(name)) = target { self.infer_definition(name); } else { // TODO infer definitions in unpacking assignment - self.infer_expression(&item.context_expr); + + // Call into the context expression inference to validate that it evaluates + // to a valid context manager. + let context_expression_ty = self.infer_expression(&item.context_expr); + self.infer_context_expression(&item.context_expr, context_expression_ty); self.infer_optional_expression(target); } } @@ -1023,18 +1031,104 @@ impl<'db> TypeInferenceBuilder<'db> { with_item: &ast::WithItem, definition: Definition<'db>, ) { - let expression = self.index.expression(&with_item.context_expr); - let result = infer_expression_types(self.db, expression); - self.extend(result); + let context_expr = self.index.expression(&with_item.context_expr); + self.extend(infer_expression_types(self.db, context_expr)); - // TODO(dhruvmanila): The correct type inference here is the return type of the __enter__ - // method of the context manager. - let context_expr_ty = self.expression_ty(&with_item.context_expr); + let target_ty = self.infer_context_expression( + &with_item.context_expr, + self.expression_ty(&with_item.context_expr), + ); self.types .expressions - .insert(target.scoped_ast_id(self.db, self.scope()), context_expr_ty); - self.add_binding(target.into(), definition, context_expr_ty); + .insert(target.scoped_ast_id(self.db, self.scope()), target_ty); + self.add_binding(target.into(), definition, target_ty); + } + + /// Infers the type of a context expression (`with expr`) and returns the target's type + /// + /// Returns [`Type::Unknown`] if the context expression doesn't implement the context manager protocol. + /// + /// ## Terminology + /// See [PEP343](https://peps.python.org/pep-0343/#standard-terminology). + fn infer_context_expression( + &mut self, + context_expression: &ast::Expr, + context_expression_ty: Type<'db>, + ) -> Type<'db> { + let context_manager_ty = context_expression_ty.to_meta_type(self.db); + + let enter_ty = context_manager_ty.member(self.db, "__enter__"); + let exit_ty = context_manager_ty.member(self.db, "__exit__"); + + // TODO: The checks here should be simplified to checking if context manager implements the `contextlib.AbstractContextManager` protocol. + if enter_ty.is_unbound() && exit_ty.is_unbound() { + self.add_diagnostic( + context_expression.into(), + "invalid-context-manager", + format_args!( + "Object of type {} cannot be used with `with` because it doesn't implement `__enter__` and `__exit__`", + context_expression_ty.display(self.db) + ), + ); + Type::Unknown + } else if enter_ty.is_unbound() { + self.add_diagnostic( + context_expression.into(), + "invalid-context-manager", + format_args!( + "Object of type {} cannot be used with `with` because it doesn't implement `__enter__`", + context_expression_ty.display(self.db) + ), + ); + Type::Unknown + } else { + let target_ty = enter_ty + .call(self.db, &[context_expression_ty]) + .return_ty_result(self.db, context_expression.into(), self) + .unwrap_or_else(|err| { + self.add_diagnostic( + context_expression.into(), + "invalid-context-manager", + format_args!(" + Object of type {context_expression} cannot be used with `with` because the method `__enter__` of type {enter_ty} is not callable", + context_expression = context_expression_ty.display(self.db), + enter_ty = enter_ty.display(self.db) + ), + ); + err.return_ty() + }); + + if exit_ty.is_unbound() { + self.add_diagnostic( + context_expression.into(), + "invalid-context-manager", + format_args!( + "Object of type {} cannot be used with `with` because it doesn't implement `__exit__`", + context_expression_ty.display(self.db) + ), + ); + } else if exit_ty + .call( + self.db, + &[context_manager_ty, Type::None, Type::None, Type::None], + ) + .return_ty_result(self.db, context_expression.into(), self) + .is_err() + { + self.add_diagnostic( + context_expression.into(), + "invalid-context-manager", + format_args!( + "Object of type {context_expression} cannot be used with `with` because the method `__exit__` of type {exit_ty} is not callable", + context_expression = context_expression_ty.display(self.db), + exit_ty = exit_ty.display(self.db), + ), + ); + } + + target_ty + } } fn infer_except_handler_definition( From 3edc006d46a81686cd27de299f6942b110a8759c Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Wed, 30 Oct 2024 18:32:04 +0100 Subject: [PATCH 2/3] "Handle" `async with` statements --- .../mdtest/declaration/async_with.md | 20 +++++++++++++ .../src/semantic_index/builder.rs | 28 +++++++++++-------- .../src/semantic_index/definition.rs | 27 +++++++++++++----- .../src/types/infer.rs | 23 ++++++++++----- 4 files changed, 73 insertions(+), 25 deletions(-) create mode 100644 crates/red_knot_python_semantic/resources/mdtest/declaration/async_with.md diff --git a/crates/red_knot_python_semantic/resources/mdtest/declaration/async_with.md b/crates/red_knot_python_semantic/resources/mdtest/declaration/async_with.md new file mode 100644 index 0000000000000..cf0fe544edea1 --- /dev/null +++ b/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 +``` diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index b8c900d6b3df4..2f4bf2c651a6a 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/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, diff --git a/crates/red_knot_python_semantic/src/semantic_index/definition.rs b/crates/red_knot_python_semantic/src/semantic_index/definition.rs index 4da4c4e6a71d6..b2ac7acee052c 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/definition.rs +++ b/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, target: AstNodeRef, + 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)] diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 9ab4cb65bd346..3b42c91bb57ee 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -476,7 +476,12 @@ impl<'db> TypeInferenceBuilder<'db> { self.infer_parameter_with_default_definition(parameter_with_default, definition); } DefinitionKind::WithItem(with_item) => { - self.infer_with_item_definition(with_item.target(), with_item.node(), definition); + self.infer_with_item_definition( + with_item.target(), + with_item.node(), + with_item.is_async(), + definition, + ); } DefinitionKind::MatchPattern(match_pattern) => { self.infer_match_pattern_definition( @@ -999,14 +1004,10 @@ impl<'db> TypeInferenceBuilder<'db> { fn infer_with_statement(&mut self, with_statement: &ast::StmtWith) { let ast::StmtWith { range: _, - is_async: _, + is_async, items, body, } = with_statement; - - // TODO: Handle async with statements. For async with statements, the value returned - // by `__enter__` must be awaited. - for item in items { let target = item.optional_vars.as_deref(); if let Some(ast::Expr::Name(name)) = target { @@ -1017,7 +1018,7 @@ impl<'db> TypeInferenceBuilder<'db> { // Call into the context expression inference to validate that it evaluates // to a valid context manager. let context_expression_ty = self.infer_expression(&item.context_expr); - self.infer_context_expression(&item.context_expr, context_expression_ty); + self.infer_context_expression(&item.context_expr, context_expression_ty, *is_async); self.infer_optional_expression(target); } } @@ -1029,6 +1030,7 @@ impl<'db> TypeInferenceBuilder<'db> { &mut self, target: &ast::ExprName, with_item: &ast::WithItem, + is_async: bool, definition: Definition<'db>, ) { let context_expr = self.index.expression(&with_item.context_expr); @@ -1037,6 +1039,7 @@ impl<'db> TypeInferenceBuilder<'db> { let target_ty = self.infer_context_expression( &with_item.context_expr, self.expression_ty(&with_item.context_expr), + is_async, ); self.types @@ -1055,7 +1058,13 @@ impl<'db> TypeInferenceBuilder<'db> { &mut self, context_expression: &ast::Expr, context_expression_ty: Type<'db>, + is_async: bool, ) -> Type<'db> { + // TODO: Handle async with statements (they use `aenter` and `aexit`) + if is_async { + return Type::Todo; + } + let context_manager_ty = context_expression_ty.to_meta_type(self.db); let enter_ty = context_manager_ty.member(self.db, "__enter__"); From 749abbda2a2b631c7b800785f18b3517e95555d6 Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Wed, 30 Oct 2024 18:40:42 +0100 Subject: [PATCH 3/3] Add test for possible unbound enter --- .../resources/mdtest/declaration/with.md | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/crates/red_knot_python_semantic/resources/mdtest/declaration/with.md b/crates/red_knot_python_semantic/resources/mdtest/declaration/with.md index 2b3bfa8f56eed..80602ec82c442 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/declaration/with.md +++ b/crates/red_knot_python_semantic/resources/mdtest/declaration/with.md @@ -120,3 +120,21 @@ context_expr = Manager1() if coinflip() else NotAContextManager() 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 +```