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
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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.
Comment on lines +5 to +6
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.


```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
```
140 changes: 140 additions & 0 deletions crates/red_knot_python_semantic/resources/mdtest/declaration/with.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,140 @@
# 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.


## 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.


```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__`"
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.

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
```
28 changes: 17 additions & 11 deletions crates/red_knot_python_semantic/src/semantic_index/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Expand Down Expand Up @@ -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,
},
);
}
Expand Down Expand Up @@ -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> {
Expand All @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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()
}
Expand Down Expand Up @@ -534,6 +542,7 @@ pub enum AssignmentKind {
pub struct WithItemDefinitionKind {
node: AstNodeRef<ast::WithItem>,
target: AstNodeRef<ast::ExprName>,
is_async: bool,
}

impl WithItemDefinitionKind {
Expand All @@ -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)]
Expand Down
2 changes: 0 additions & 2 deletions crates/red_knot_python_semantic/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.

///
/// Returns `None` if `self` is not a callable type.
MichaReiser marked this conversation as resolved.
Show resolved Hide resolved
#[must_use]
fn call(self, db: &'db dyn Db, arg_types: &[Type<'db>]) -> CallOutcome<'db> {
match self {
Expand Down
Loading
Loading