Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from

Conversation

MichaReiser
Copy link
Member

Summary

This PR implements, hopefully correct, inference for the target types of with items (sync with statements only).

The context expression specifies a context manager and the target type must be inferred to the return type of the context manager's __enter__ method.
That's what this PR implements.

The ad-hoc handling of protocols is awkward and leads to a lot of repetitive code but I decided not to spent too much time on it because
asserting that the context expression correctly implements the context manager protocol is going to change with proper protocol support.

Test Plan

Added tests

@MichaReiser MichaReiser added the red-knot Multi-file analysis & type inference label Oct 30, 2024
Copy link
Contributor

github-actions bot commented Oct 30, 2024

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

AlexWaygood
AlexWaygood previously approved these changes Oct 30, 2024
Copy link
Member

@AlexWaygood AlexWaygood left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!

I'm not sure this properly handles the case where e.g. __enter__ might or might not be defined, e.g.

def coinflip() -> bool:
    return True

class Manager:
    if coinflip():
        def __enter__(self) -> Target:
            return Target()

    def __exit__(self, *args): ...

I'm also not sure it's essential that we tackle that now for an initial implementation, but it could be worth at least adding a TODO around it

crates/red_knot_python_semantic/src/types.rs Show resolved Hide resolved
crates/red_knot_python_semantic/src/types/infer.rs Outdated Show resolved Hide resolved
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.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't actually know if simplification here is possible because simply checking if context_manager_ty is a subtype of AbstractContextManager won't give us the return type of the __enter__ method. I think we do actually need to manually call the __enter__ method somehow, and once we've done that we then have to manually consider the case where __enter__ isn't actually callable... it's hard for me to see what would actually be simpler once we have generalised support for typing.Protocol

Well, one thing we could do is to just check that the __exit__ method is consistent with the __exit__ method on AbstractContextManager, since __exit__ method definitions can get pretty darn complicated. But we're not checking the call signature of __exit__ at all right now, so this would be "fending off future complexity" rather than reducing any of the complexity that's already here

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think what would become reusable is testing if the context manager has the right shape. Once we know that, we can call enter (which should never fail at this point)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think maybe what Alex is implying is that it could be kind of wasteful/unnecessary to check __enter__ against the protocol, when simply calling it and seeing if the call succeeds or fails is actually sufficient. But this is a question that we can leave for later, I think the comment is reasonable as-is. We can always change our mind when actually implementing a TODO if it turns out not to make sense to to do it that way.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

whatever the case, this is a very long line ;)

Suggested change
// TODO: The checks here should be simplified to checking if context manager implements the `contextlib.AbstractContextManager` protocol.
// TODO: The checks here should be simplified to checking
// if context manager implements the `contextlib.AbstractContextManager` protocol.

context_expression.into(),
"invalid-context-manager",
format_args!(
"Object of type {} cannot be used with `with` because it doesn't implement `__enter__` and `__exit__`",
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The detail in these error messages is fantastic! Long-term, I'm not sure we'll want all this information in one long sentence like this, but I think this is fine for now

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree

@MichaReiser
Copy link
Member Author

I'm not sure this properly handles the case where e.g. enter might or might not be defined, e.g.

Lol, python allows if statements in class bodies

Copy link
Member

@AlexWaygood AlexWaygood left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, it's okay to skip implementing the logic for async with statements yet, and just infer Todo for them. But this PR introduces false positives for async with statements, and I don't think that's okay.

On this snippet:

class Foo:
    async def __aenter__(self):
        return "foo"

    async def __aexit__(self, *args):
        pass

async def foo():
    async with Foo() as x:
        reveal_type(x)

red-knot reports with your PR branch:

ERROR /Users/alexw/dev/experiment/foo.py:9:16: Object of type Foo cannot be used with `with` because it doesn't implement `__enter__` and `__exit__`
ERROR /Users/alexw/dev/experiment/foo.py:10:9: `reveal_type` used without importing it; this is allowed for debugging convenience but will fail at runtime
ERROR /Users/alexw/dev/experiment/foo.py:10:9: Revealed type is `Unknown`

I think you might need to do something like this to avoid adding false positives for async with statements in this PR:

--- 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,12 @@ 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 { is_async: *is_async, item });
                         self.visit_expr(optional_vars);
                         self.pop_assignment();
                     }
@@ -1011,11 +1011,12 @@ where
                                 },
                             );
                         }
-                        Some(CurrentAssignment::WithItem(with_item)) => {
+                        Some(CurrentAssignment::WithItem{is_async, item}) => {
                             self.add_definition(
                                 symbol,
                                 WithItemDefinitionNodeRef {
-                                    node: with_item,
+                                    is_async,
+                                    node: item,
                                     target: name_node,
                                 },
                             );
@@ -1232,7 +1233,10 @@ enum CurrentAssignment<'a> {
         node: &'a ast::Comprehension,
         first: bool,
     },
-    WithItem(&'a ast::WithItem),
+    WithItem {
+        is_async: bool,
+        item: &'a ast::WithItem,
+    },
 }
 
 impl<'a> From<&'a ast::StmtAnnAssign> for CurrentAssignment<'a> {
@@ -1259,12 +1263,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 4da4c4e6a..f1826d1de 100644
--- a/crates/red_knot_python_semantic/src/semantic_index/definition.rs
+++ b/crates/red_knot_python_semantic/src/semantic_index/definition.rs
@@ -174,6 +174,7 @@ pub(crate) struct AssignmentDefinitionNodeRef<'a> {
 
 #[derive(Copy, Clone, Debug)]
 pub(crate) struct WithItemDefinitionNodeRef<'a> {
+    pub(crate) is_async: bool,
     pub(crate) node: &'a ast::WithItem,
     pub(crate) target: &'a ast::ExprName,
 }
@@ -277,8 +278,9 @@ impl DefinitionNodeRef<'_> {
                     DefinitionKind::ParameterWithDefault(AstNodeRef::new(parsed, parameter))
                 }
             },
-            DefinitionNodeRef::WithItem(WithItemDefinitionNodeRef { node, target }) => {
+            DefinitionNodeRef::WithItem(WithItemDefinitionNodeRef { is_async, node, target }) => {
                 DefinitionKind::WithItem(WithItemDefinitionKind {
+                    is_async,
                     node: AstNodeRef::new(parsed.clone(), node),
                     target: AstNodeRef::new(parsed, target),
                 })
@@ -329,7 +331,7 @@ impl DefinitionNodeRef<'_> {
                 ast::AnyParameterRef::Variadic(parameter) => parameter.into(),
                 ast::AnyParameterRef::NonVariadic(parameter) => parameter.into(),
             },
-            Self::WithItem(WithItemDefinitionNodeRef { node: _, target }) => target.into(),
+            Self::WithItem(WithItemDefinitionNodeRef { is_async: _, node: _, target }) => target.into(),
             Self::MatchPattern(MatchPatternDefinitionNodeRef { identifier, .. }) => {
                 identifier.into()
             }
@@ -532,11 +534,16 @@ pub enum AssignmentKind {
 
 #[derive(Clone, Debug)]
 pub struct WithItemDefinitionKind {
+    is_async: bool,
     node: AstNodeRef<ast::WithItem>,
     target: AstNodeRef<ast::ExprName>,
 }
 
 impl WithItemDefinitionKind {
+    pub(crate) fn is_async(&self) -> bool {
+        self.is_async
+    }
+
     pub(crate) fn node(&self) -> &ast::WithItem {
         self.node.node()
     }
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index 9ab4cb65b..6550f2ee2 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.is_async(),
+                    with_item.target(),
+                    with_item.node(),
+                    definition,
+                );
             }
             DefinitionKind::MatchPattern(match_pattern) => {
                 self.infer_match_pattern_definition(
@@ -999,7 +1004,7 @@ 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;
@@ -1017,7 +1022,9 @@ 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);
+                if !is_async {
+                    self.infer_context_expression(&item.context_expr, context_expression_ty);
+                }
                 self.infer_optional_expression(target);
             }
         }
@@ -1027,6 +1034,7 @@ impl<'db> TypeInferenceBuilder<'db> {
 
     fn infer_with_item_definition(
         &mut self,
+        is_async: bool,
         target: &ast::ExprName,
         with_item: &ast::WithItem,
         definition: Definition<'db>,
@@ -1034,10 +1042,14 @@ impl<'db> TypeInferenceBuilder<'db> {
         let context_expr = self.index.expression(&with_item.context_expr);
         self.extend(infer_expression_types(self.db, context_expr));
 
-        let target_ty = self.infer_context_expression(
-            &with_item.context_expr,
-            self.expression_ty(&with_item.context_expr),
-        );
+        let target_ty = if is_async {
+            Type::Todo
+        } else {
+            self.infer_context_expression(
+                &with_item.context_expr,
+                self.expression_ty(&with_item.context_expr),
+            )
+        };
 
         self.types
             .expressions

@AlexWaygood AlexWaygood dismissed their stale review October 30, 2024 16:14

introduces false positives for async with

@MichaReiser
Copy link
Member Author

I'm not sure this properly handles the case where e.g. enter might or might not be defined, e.g.

nice catch. I added a test with a todo. I want to wait with handling this until we removed unbound.

Copy link
Member

@AlexWaygood AlexWaygood left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! This LGTM now.

Comment on lines +5 to +6
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.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: quite long lines here. Also, we haven't renamed it to Knot yet ;)

Suggested change
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.
The type of the target variable in a `with` statement should be the return type
from the context manager's `__aenter__` method. However, we don't support `async with`
statements yet. This test asserts that it doesn't emit any context manager-related errors.


## Basic `with` statement

The type of the target variable in a `with` statement is the return type from the context manager's `__enter__` method.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
The type of the target variable in a `with` statement is the return type from the context manager's `__enter__` method.
The type of the target variable in a `with` statement is the return type
from the context manager's `__enter__` method.

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.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

whatever the case, this is a very long line ;)

Suggested change
// TODO: The checks here should be simplified to checking if context manager implements the `contextlib.AbstractContextManager` protocol.
// TODO: The checks here should be simplified to checking
// if context manager implements the `contextlib.AbstractContextManager` protocol.

Copy link
Contributor

@carljm carljm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is really well done! Super clear, thorough, and well-researched.

@@ -0,0 +1,122 @@
# With declarations
Copy link
Contributor

@carljm carljm Oct 30, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit about this title, and the location of this file: with statements create Definitions, which are Bindings, but are not Declarations. (A declaration is only an annotated assignment, or an annotated function parameter, or a function or class definition.) So I think this title should be just With statements.

Regarding file location, we don't have a general directory for all kinds of definitions/bindings, so I think with/ could be a top-level directory if/when we have multiple files to put in there; for now a top level with.md would be fine, I think.

```py
class Manager: ...

# error: [invalid-context-manager] "Object of type Manager cannot be used with `with` because it doesn't implement `__enter__` and `__exit__`"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The diagnostic messages here and below are excellent! Super clear and informative.

@@ -1010,8 +1010,6 @@ impl<'db> Type<'db> {
}

/// Return the type resulting from calling an object of this type.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While we're editing this comment anyway, this part isn't accurate anymore either. I think it could be edited as below. Arguably that doesn't give much information that isn't obvious, but I think it's not 100% obvious that Type::call means "what happens if I call an object of this type", so I'd probably still favor keeping the short doc comment.

We could also do a much longer doc comment here or elsewhere that talks about how to actually correctly use CallOutcome, but that's definitely out of scope for this PR :)

Suggested change
/// Return the type resulting from calling an object of this type.
/// Return the outcome of calling an object of this type.

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.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think maybe what Alex is implying is that it could be kind of wasteful/unnecessary to check __enter__ against the protocol, when simply calling it and seeing if the call succeeds or fails is actually sufficient. But this is a question that we can leave for later, I think the comment is reasonable as-is. We can always change our mind when actually implementing a TODO if it turns out not to make sense to to do it that way.

self.db,
&[context_manager_ty, Type::None, Type::None, Type::None],
)
.return_ty_result(self.db, context_expression.into(), self)
Copy link
Contributor

@carljm carljm Oct 30, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The return type here will matter to us later, because it will determine whether exceptions are suppressed by this with block. We could add a TODO for that but I don't think it's important that we do.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
red-knot Multi-file analysis & type inference
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants