Skip to content

Commit

Permalink
Understand type[A | B] special form in annotations (#14830)
Browse files Browse the repository at this point in the history
resolves #14703

I decided to use recursion to get the type, so if anything is added to
the single element inference it will be applied for the union.
Also added this
[change](#14703 (comment))
in this PR since it was easy.

---------

Co-authored-by: Carl Meyer <[email protected]>
  • Loading branch information
Glyphack and carljm authored Dec 7, 2024
1 parent d340134 commit 269e47b
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 1 deletion.
28 changes: 28 additions & 0 deletions crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md
Original file line number Diff line number Diff line change
Expand Up @@ -87,3 +87,31 @@ reveal_type(f()) # revealed: @Todo(unsupported type[X] special form)
```py path=a/b.py
class C: ...
```

## Union of classes

```py
class BasicUser: ...
class ProUser: ...

class A:
class B:
class C: ...

def get_user() -> type[BasicUser | ProUser | A.B.C]:
return BasicUser

# revealed: type[BasicUser] | type[ProUser] | type[C]
reveal_type(get_user())
```

## Illegal parameters

```py
class A: ...
class B: ...

# error: [invalid-type-form]
def get_user() -> type[A, B]:
return A
```
23 changes: 22 additions & 1 deletion crates/red_knot_python_semantic/src/types/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4663,7 +4663,28 @@ impl<'db> TypeInferenceBuilder<'db> {
todo_type!("unsupported type[X] special form")
}
}
// TODO: unions, subscripts, etc.
ast::Expr::BinOp(binary) if binary.op == ast::Operator::BitOr => {
let union_ty = UnionType::from_elements(
self.db,
[
self.infer_subclass_of_type_expression(&binary.left),
self.infer_subclass_of_type_expression(&binary.right),
],
);
self.store_expression_type(slice, union_ty);

union_ty
}
ast::Expr::Tuple(_) => {
self.infer_type_expression(slice);
self.diagnostics.add(
slice.into(),
"invalid-type-form",
format_args!("type[...] must have exactly one type argument"),
);
Type::Unknown
}
// TODO: subscripts, etc.
_ => {
self.infer_type_expression(slice);
todo_type!("unsupported type[X] special form")
Expand Down

0 comments on commit 269e47b

Please sign in to comment.