Skip to content

Commit aecdb8c

Browse files
Glyphackcarljm
andauthored
[red-knot] support typing.Union in type annotations (#14499)
Fix #14498 ## Summary This PR adds `typing.Union` support ## Test Plan I created new tests in mdtest. --------- Co-authored-by: Carl Meyer <[email protected]>
1 parent 3c52d2d commit aecdb8c

File tree

4 files changed

+78
-1
lines changed

4 files changed

+78
-1
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
# Union
2+
3+
## Annotation
4+
5+
`typing.Union` can be used to construct union types same as `|` operator.
6+
7+
```py
8+
from typing import Union
9+
10+
a: Union[int, str]
11+
a1: Union[int, bool]
12+
a2: Union[int, Union[float, str]]
13+
a3: Union[int, None]
14+
a4: Union[Union[float, str]]
15+
a5: Union[int]
16+
a6: Union[()]
17+
18+
def f():
19+
# revealed: int | str
20+
reveal_type(a)
21+
# Since bool is a subtype of int we simplify to int here. But we do allow assigning boolean values (see below).
22+
# revealed: int
23+
reveal_type(a1)
24+
# revealed: int | float | str
25+
reveal_type(a2)
26+
# revealed: int | None
27+
reveal_type(a3)
28+
# revealed: float | str
29+
reveal_type(a4)
30+
# revealed: int
31+
reveal_type(a5)
32+
# revealed: Never
33+
reveal_type(a6)
34+
```
35+
36+
## Assignment
37+
38+
```py
39+
from typing import Union
40+
41+
a: Union[int, str]
42+
a = 1
43+
a = ""
44+
a1: Union[int, bool]
45+
a1 = 1
46+
a1 = True
47+
# error: [invalid-assignment] "Object of type `Literal[b""]` is not assignable to `int | str`"
48+
a = b""
49+
```
50+
51+
## Typing Extensions
52+
53+
```py
54+
from typing_extensions import Union
55+
56+
a: Union[int, str]
57+
58+
def f():
59+
# revealed: int | str
60+
reveal_type(a)
61+
```

crates/red_knot_python_semantic/src/types.rs

+9-1
Original file line numberDiff line numberDiff line change
@@ -1809,6 +1809,8 @@ pub enum KnownInstanceType<'db> {
18091809
Literal,
18101810
/// The symbol `typing.Optional` (which can also be found as `typing_extensions.Optional`)
18111811
Optional,
1812+
/// The symbol `typing.Union` (which can also be found as `typing_extensions.Union`)
1813+
Union,
18121814
/// A single instance of `typing.TypeVar`
18131815
TypeVar(TypeVarInstance<'db>),
18141816
// TODO: fill this enum out with more special forms, etc.
@@ -1819,14 +1821,17 @@ impl<'db> KnownInstanceType<'db> {
18191821
match self {
18201822
KnownInstanceType::Literal => "Literal",
18211823
KnownInstanceType::Optional => "Optional",
1824+
KnownInstanceType::Union => "Union",
18221825
KnownInstanceType::TypeVar(_) => "TypeVar",
18231826
}
18241827
}
18251828

18261829
/// Evaluate the known instance in boolean context
18271830
pub const fn bool(self) -> Truthiness {
18281831
match self {
1829-
Self::Literal | Self::Optional | Self::TypeVar(_) => Truthiness::AlwaysTrue,
1832+
Self::Literal | Self::Optional | Self::TypeVar(_) | Self::Union => {
1833+
Truthiness::AlwaysTrue
1834+
}
18301835
}
18311836
}
18321837

@@ -1835,6 +1840,7 @@ impl<'db> KnownInstanceType<'db> {
18351840
match self {
18361841
Self::Literal => "typing.Literal",
18371842
Self::Optional => "typing.Optional",
1843+
Self::Union => "typing.Union",
18381844
Self::TypeVar(typevar) => typevar.name(db),
18391845
}
18401846
}
@@ -1844,6 +1850,7 @@ impl<'db> KnownInstanceType<'db> {
18441850
match self {
18451851
Self::Literal => KnownClass::SpecialForm,
18461852
Self::Optional => KnownClass::SpecialForm,
1853+
Self::Union => KnownClass::SpecialForm,
18471854
Self::TypeVar(_) => KnownClass::TypeVar,
18481855
}
18491856
}
@@ -1864,6 +1871,7 @@ impl<'db> KnownInstanceType<'db> {
18641871
match (module.name().as_str(), instance_name) {
18651872
("typing" | "typing_extensions", "Literal") => Some(Self::Literal),
18661873
("typing" | "typing_extensions", "Optional") => Some(Self::Optional),
1874+
("typing" | "typing_extensions", "Union") => Some(Self::Union),
18671875
_ => None,
18681876
}
18691877
}

crates/red_knot_python_semantic/src/types/infer.rs

+7
Original file line numberDiff line numberDiff line change
@@ -4567,6 +4567,13 @@ impl<'db> TypeInferenceBuilder<'db> {
45674567
let param_type = self.infer_type_expression(parameters);
45684568
UnionType::from_elements(self.db, [param_type, Type::none(self.db)])
45694569
}
4570+
KnownInstanceType::Union => match parameters {
4571+
ast::Expr::Tuple(t) => UnionType::from_elements(
4572+
self.db,
4573+
t.iter().map(|elt| self.infer_type_expression(elt)),
4574+
),
4575+
_ => self.infer_type_expression(parameters),
4576+
},
45704577
KnownInstanceType::TypeVar(_) => Type::Todo,
45714578
}
45724579
}

crates/red_knot_python_semantic/src/types/mro.rs

+1
Original file line numberDiff line numberDiff line change
@@ -373,6 +373,7 @@ impl<'db> ClassBase<'db> {
373373
Type::KnownInstance(known_instance) => match known_instance {
374374
KnownInstanceType::TypeVar(_)
375375
| KnownInstanceType::Literal
376+
| KnownInstanceType::Union
376377
| KnownInstanceType::Optional => None,
377378
},
378379
}

0 commit comments

Comments
 (0)