Enable auto-return-type involving Optional and Union annotations (#8885)

## Summary

Previously, this was only supported for Python 3.10 and later, since we
always use the PEP 604-style unions.
This commit is contained in:
Charlie Marsh 2023-11-28 18:35:55 -08:00 committed by GitHub
parent ec7456bac0
commit 6435e4e4aa
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 526 additions and 69 deletions

View file

@ -1293,6 +1293,50 @@ pub fn pep_604_union(elts: &[Expr]) -> Expr {
}
}
pub fn typing_optional(elt: Expr, binding: String) -> Expr {
Expr::Subscript(ast::ExprSubscript {
value: Box::new(Expr::Name(ast::ExprName {
id: binding,
range: TextRange::default(),
ctx: ExprContext::Load,
})),
slice: Box::new(elt),
ctx: ExprContext::Load,
range: TextRange::default(),
})
}
pub fn typing_union(elts: &[Expr], binding: String) -> Expr {
fn tuple(elts: &[Expr]) -> Expr {
match elts {
[] => Expr::Tuple(ast::ExprTuple {
elts: vec![],
ctx: ExprContext::Load,
range: TextRange::default(),
}),
[Expr::Tuple(ast::ExprTuple { elts, .. })] => pep_604_union(elts),
[elt] => elt.clone(),
[rest @ .., elt] => Expr::BinOp(ast::ExprBinOp {
left: Box::new(tuple(rest)),
op: Operator::BitOr,
right: Box::new(elt.clone()),
range: TextRange::default(),
}),
}
}
Expr::Subscript(ast::ExprSubscript {
value: Box::new(Expr::Name(ast::ExprName {
id: binding,
range: TextRange::default(),
ctx: ExprContext::Load,
})),
slice: Box::new(tuple(elts)),
ctx: ExprContext::Load,
range: TextRange::default(),
})
}
#[cfg(test)]
mod tests {
use std::borrow::Cow;