[red-knot] support typing.Union in type annotations (#14499)
Some checks are pending
CI / cargo clippy (push) Blocked by required conditions
CI / cargo test (linux) (push) Blocked by required conditions
CI / cargo test (windows) (push) Blocked by required conditions
CI / cargo test (wasm) (push) Blocked by required conditions
CI / cargo build (release) (push) Blocked by required conditions
CI / cargo shear (push) Blocked by required conditions
CI / python package (push) Waiting to run
CI / pre-commit (push) Waiting to run
CI / mkdocs (push) Waiting to run
CI / formatter instabilities and black similarity (push) Blocked by required conditions
CI / Determine changes (push) Waiting to run
CI / cargo fmt (push) Waiting to run
CI / cargo build (msrv) (push) Blocked by required conditions
CI / cargo fuzz build (push) Blocked by required conditions
CI / fuzz parser (push) Blocked by required conditions
CI / test scripts (push) Blocked by required conditions
CI / ecosystem (push) Blocked by required conditions
CI / test ruff-lsp (push) Blocked by required conditions
CI / benchmarks (push) Blocked by required conditions

Fix #14498

## Summary

This PR adds `typing.Union` support

## Test Plan

I created new tests in mdtest.

---------

Co-authored-by: Carl Meyer <carl@astral.sh>
This commit is contained in:
Shaygan Hooshyari 2024-11-20 22:55:33 +01:00 committed by GitHub
parent 3c52d2d1bd
commit aecdb8c144
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 78 additions and 1 deletions

View file

@ -0,0 +1,61 @@
# Union
## Annotation
`typing.Union` can be used to construct union types same as `|` operator.
```py
from typing import Union
a: Union[int, str]
a1: Union[int, bool]
a2: Union[int, Union[float, str]]
a3: Union[int, None]
a4: Union[Union[float, str]]
a5: Union[int]
a6: Union[()]
def f():
# revealed: int | str
reveal_type(a)
# Since bool is a subtype of int we simplify to int here. But we do allow assigning boolean values (see below).
# revealed: int
reveal_type(a1)
# revealed: int | float | str
reveal_type(a2)
# revealed: int | None
reveal_type(a3)
# revealed: float | str
reveal_type(a4)
# revealed: int
reveal_type(a5)
# revealed: Never
reveal_type(a6)
```
## Assignment
```py
from typing import Union
a: Union[int, str]
a = 1
a = ""
a1: Union[int, bool]
a1 = 1
a1 = True
# error: [invalid-assignment] "Object of type `Literal[b""]` is not assignable to `int | str`"
a = b""
```
## Typing Extensions
```py
from typing_extensions import Union
a: Union[int, str]
def f():
# revealed: int | str
reveal_type(a)
```

View file

@ -1809,6 +1809,8 @@ pub enum KnownInstanceType<'db> {
Literal,
/// The symbol `typing.Optional` (which can also be found as `typing_extensions.Optional`)
Optional,
/// The symbol `typing.Union` (which can also be found as `typing_extensions.Union`)
Union,
/// A single instance of `typing.TypeVar`
TypeVar(TypeVarInstance<'db>),
// TODO: fill this enum out with more special forms, etc.
@ -1819,6 +1821,7 @@ impl<'db> KnownInstanceType<'db> {
match self {
KnownInstanceType::Literal => "Literal",
KnownInstanceType::Optional => "Optional",
KnownInstanceType::Union => "Union",
KnownInstanceType::TypeVar(_) => "TypeVar",
}
}
@ -1826,7 +1829,9 @@ impl<'db> KnownInstanceType<'db> {
/// Evaluate the known instance in boolean context
pub const fn bool(self) -> Truthiness {
match self {
Self::Literal | Self::Optional | Self::TypeVar(_) => Truthiness::AlwaysTrue,
Self::Literal | Self::Optional | Self::TypeVar(_) | Self::Union => {
Truthiness::AlwaysTrue
}
}
}
@ -1835,6 +1840,7 @@ impl<'db> KnownInstanceType<'db> {
match self {
Self::Literal => "typing.Literal",
Self::Optional => "typing.Optional",
Self::Union => "typing.Union",
Self::TypeVar(typevar) => typevar.name(db),
}
}
@ -1844,6 +1850,7 @@ impl<'db> KnownInstanceType<'db> {
match self {
Self::Literal => KnownClass::SpecialForm,
Self::Optional => KnownClass::SpecialForm,
Self::Union => KnownClass::SpecialForm,
Self::TypeVar(_) => KnownClass::TypeVar,
}
}
@ -1864,6 +1871,7 @@ impl<'db> KnownInstanceType<'db> {
match (module.name().as_str(), instance_name) {
("typing" | "typing_extensions", "Literal") => Some(Self::Literal),
("typing" | "typing_extensions", "Optional") => Some(Self::Optional),
("typing" | "typing_extensions", "Union") => Some(Self::Union),
_ => None,
}
}

View file

@ -4567,6 +4567,13 @@ impl<'db> TypeInferenceBuilder<'db> {
let param_type = self.infer_type_expression(parameters);
UnionType::from_elements(self.db, [param_type, Type::none(self.db)])
}
KnownInstanceType::Union => match parameters {
ast::Expr::Tuple(t) => UnionType::from_elements(
self.db,
t.iter().map(|elt| self.infer_type_expression(elt)),
),
_ => self.infer_type_expression(parameters),
},
KnownInstanceType::TypeVar(_) => Type::Todo,
}
}

View file

@ -373,6 +373,7 @@ impl<'db> ClassBase<'db> {
Type::KnownInstance(known_instance) => match known_instance {
KnownInstanceType::TypeVar(_)
| KnownInstanceType::Literal
| KnownInstanceType::Union
| KnownInstanceType::Optional => None,
},
}