Skip to content

Commit d8bb65d

Browse files
asukaminato0721asukaminato0721
committed
fix?
1 parent 47a64ef commit d8bb65d

File tree

3 files changed

+140
-51
lines changed

3 files changed

+140
-51
lines changed

‎pyrefly/lib/alt/expr.rs‎

Lines changed: 44 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ use ruff_python_ast::DictItem;
3030
use ruff_python_ast::Expr;
3131
use ruff_python_ast::ExprCall;
3232
use ruff_python_ast::ExprGenerator;
33+
use ruff_python_ast::ExprList;
3334
use ruff_python_ast::ExprNumberLiteral;
3435
use ruff_python_ast::ExprSlice;
3536
use ruff_python_ast::ExprStarred;
@@ -419,29 +420,20 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
419420
}
420421
Expr::Tuple(x) => self.tuple_infer(x, hint, errors),
421422
Expr::List(x) => {
422-
let elt_hint = hint.and_then(|ty| self.decompose_list(ty));
423-
if x.is_empty() {
424-
let elem_ty = elt_hint.map_or_else(
425-
|| {
426-
if !self.solver().infer_with_first_use {
427-
self.error(
428-
errors,
429-
x.range(),
430-
ErrorInfo::Kind(ErrorKind::ImplicitAny),
431-
"This expression is implicitly inferred to be `list[Any]`. Please provide an explicit type annotation.".to_owned(),
432-
);
433-
Type::any_implicit()
434-
} else {
435-
self.solver().fresh_contained(self.uniques).to_type()
436-
}
437-
},
438-
|hint| hint.to_type(),
439-
);
440-
self.stdlib.list(elem_ty).to_type()
441-
} else {
442-
let elem_tys = self.elts_infer(&x.elts, elt_hint, errors);
443-
self.stdlib.list(self.unions(elem_tys)).to_type()
423+
if let Some(hint_ref) = hint.as_ref()
424+
&& let Type::Union(options) = hint_ref.ty()
425+
{
426+
for option in options {
427+
let branch_hint =
428+
self.decompose_list(HintRef::new(option, hint_ref.errors()));
429+
let ty = self.list_with_hint(x, branch_hint, errors);
430+
if self.is_subset_eq(&ty, option) {
431+
return ty;
432+
}
433+
}
444434
}
435+
let elt_hint = hint.and_then(|ty| self.decompose_list(ty));
436+
self.list_with_hint(x, elt_hint, errors)
445437
}
446438
Expr::Dict(x) => self.dict_infer(&x.items, hint, x.range, errors),
447439
Expr::Set(x) => {
@@ -1815,6 +1807,36 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
18151807
})
18161808
}
18171809

1810+
fn list_with_hint(
1811+
&self,
1812+
x: &ExprList,
1813+
elt_hint: Option<Hint>,
1814+
errors: &ErrorCollector,
1815+
) -> Type {
1816+
if x.is_empty() {
1817+
let elem_ty = elt_hint.map_or_else(
1818+
|| {
1819+
if !self.solver().infer_with_first_use {
1820+
self.error(
1821+
errors,
1822+
x.range(),
1823+
ErrorInfo::Kind(ErrorKind::ImplicitAny),
1824+
"This expression is implicitly inferred to be `list[Any]`. Please provide an explicit type annotation.".to_owned(),
1825+
);
1826+
Type::any_implicit()
1827+
} else {
1828+
self.solver().fresh_contained(self.uniques).to_type()
1829+
}
1830+
},
1831+
|hint| hint.to_type(),
1832+
);
1833+
self.stdlib.list(elem_ty).to_type()
1834+
} else {
1835+
let elem_tys = self.elts_infer(&x.elts, elt_hint, errors);
1836+
self.stdlib.list(self.unions(elem_tys)).to_type()
1837+
}
1838+
}
1839+
18181840
fn intercept_typing_self_use(&self, x: &Expr) -> Option<TypeInfo> {
18191841
match x {
18201842
Expr::Name(..) | Expr::Attribute(..) => {

‎pyrefly/lib/alt/unwrap.rs‎

Lines changed: 95 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,54 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
119119
}
120120
}
121121

122+
fn collect_var_from_hint<F>(&self, ty: &Type, make: &F) -> Option<Vec<Type>>
123+
where
124+
F: Fn(Var) -> Type,
125+
{
126+
match ty {
127+
Type::Union(tys) => {
128+
let mut collected = Vec::new();
129+
let mut matched = false;
130+
for branch in tys {
131+
if let Some(mut branch_res) = self.collect_var_from_hint(branch, make) {
132+
matched = true;
133+
collected.append(&mut branch_res);
134+
}
135+
}
136+
if matched { Some(collected) } else { None }
137+
}
138+
_ => {
139+
let var = self.fresh_var();
140+
let target = make(var);
141+
if self.is_subset_eq(&target, ty) {
142+
match self.resolve_var_opt(ty, var) {
143+
Some(resolved) => Some(vec![resolved]),
144+
None => Some(Vec::new()),
145+
}
146+
} else {
147+
None
148+
}
149+
}
150+
}
151+
}
152+
153+
fn hint_from_types<'b>(
154+
&self,
155+
mut types: Vec<Type>,
156+
hint: &HintRef<'b, '_>,
157+
) -> Option<Hint<'b>> {
158+
if types.is_empty() {
159+
None
160+
} else {
161+
let ty = if types.len() == 1 {
162+
types.pop().unwrap()
163+
} else {
164+
self.unions(types)
165+
};
166+
Some(hint.map_ty(|_| ty))
167+
}
168+
}
169+
122170
pub fn unwrap_mapping(&self, ty: &Type) -> Option<(Type, Type)> {
123171
let key = self.fresh_var();
124172
let value = self.fresh_var();
@@ -224,45 +272,65 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
224272
&self,
225273
hint: HintRef<'b, '_>,
226274
) -> (Option<Hint<'b>>, Option<Hint<'b>>) {
227-
let key = self.fresh_var();
228-
let value = self.fresh_var();
229-
let dict_type = self.stdlib.dict(key.to_type(), value.to_type()).to_type();
230-
if self.is_subset_eq(&dict_type, hint.ty()) {
231-
let key = hint.map_ty_opt(|ty| self.resolve_var_opt(ty, key));
232-
let value = hint.map_ty_opt(|ty| self.resolve_var_opt(ty, value));
233-
(key, value)
234-
} else {
235-
(None, None)
275+
let mut key_types = Vec::new();
276+
let mut value_types = Vec::new();
277+
let mut matched = false;
278+
279+
// Helper to process a single target type and accumulate results.
280+
let mut consider = |ty: &Type| {
281+
let key = self.fresh_var();
282+
let value = self.fresh_var();
283+
let dict_type = self.stdlib.dict(key.to_type(), value.to_type()).to_type();
284+
if self.is_subset_eq(&dict_type, ty) {
285+
matched = true;
286+
if let Some(key_ty) = self.resolve_var_opt(ty, key) {
287+
key_types.push(key_ty);
288+
}
289+
if let Some(value_ty) = self.resolve_var_opt(ty, value) {
290+
value_types.push(value_ty);
291+
}
292+
}
293+
};
294+
295+
match hint.ty() {
296+
Type::Union(branches) => {
297+
for branch in branches {
298+
consider(branch);
299+
}
300+
}
301+
ty => consider(ty),
236302
}
303+
304+
if !matched {
305+
return (None, None);
306+
}
307+
308+
let key = self.hint_from_types(key_types, &hint);
309+
let value = self.hint_from_types(value_types, &hint);
310+
(key, value)
237311
}
238312

239313
pub fn decompose_set<'b>(&self, hint: HintRef<'b, '_>) -> Option<Hint<'b>> {
240-
let elem = self.fresh_var();
241-
let set_type = self.stdlib.set(elem.to_type()).to_type();
242-
if self.is_subset_eq(&set_type, hint.ty()) {
243-
hint.map_ty_opt(|ty| self.resolve_var_opt(ty, elem))
244-
} else {
245-
None
314+
let make = |var: Var| self.stdlib.set(var.to_type()).to_type();
315+
match self.collect_var_from_hint(hint.ty(), &make) {
316+
Some(tys) => self.hint_from_types(tys, &hint),
317+
None => None,
246318
}
247319
}
248320

249321
pub fn decompose_list<'b>(&self, hint: HintRef<'b, '_>) -> Option<Hint<'b>> {
250-
let elem = self.fresh_var();
251-
let list_type = self.stdlib.list(elem.to_type()).to_type();
252-
if self.is_subset_eq(&list_type, hint.ty()) {
253-
hint.map_ty_opt(|ty| self.resolve_var_opt(ty, elem))
254-
} else {
255-
None
322+
let make = |var: Var| self.stdlib.list(var.to_type()).to_type();
323+
match self.collect_var_from_hint(hint.ty(), &make) {
324+
Some(tys) => self.hint_from_types(tys, &hint),
325+
None => None,
256326
}
257327
}
258328

259329
pub fn decompose_tuple<'b>(&self, hint: HintRef<'b, '_>) -> Option<Hint<'b>> {
260-
let elem = self.fresh_var();
261-
let tuple_type = self.stdlib.tuple(elem.to_type()).to_type();
262-
if self.is_subset_eq(&tuple_type, hint.ty()) {
263-
hint.map_ty_opt(|ty| self.resolve_var_opt(ty, elem))
264-
} else {
265-
None
330+
let make = |var: Var| self.stdlib.tuple(var.to_type()).to_type();
331+
match self.collect_var_from_hint(hint.ty(), &make) {
332+
Some(tys) => self.hint_from_types(tys, &hint),
333+
None => None,
266334
}
267335
}
268336

‎pyrefly/lib/test/contextual.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -102,15 +102,14 @@ kwarg(xs=[B()], ys=[B()])
102102
);
103103

104104
testcase!(
105-
bug = "Both assignments should be allowed. When decomposing the contextual hint, we eagerly resolve vars to the 'first' branch of the union. Note: due to the union's sorted representation, the first branch is not necessarily the first in source order.",
106105
test_contextual_typing_against_unions,
107106
r#"
108107
class A: ...
109108
class B: ...
110109
class B2(B): ...
111110
class C: ...
112111
113-
x: list[A] | list[B] = [B2()] # E: `list[B2]` is not assignable to `list[A] | list[B]`
112+
x: list[A] | list[B] = [B2()]
114113
y: list[B] | list[C] = [B2()]
115114
"#,
116115
);

0 commit comments

Comments
 (0)