Skip to content

Commit c95c36b

Browse files
asukaminato0721asukaminato0721
committed
fix?
1 parent 47b5af5 commit c95c36b

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;
@@ -418,29 +419,20 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
418419
}
419420
Expr::Tuple(x) => self.tuple_infer(x, hint, errors),
420421
Expr::List(x) => {
421-
let elt_hint = hint.and_then(|ty| self.decompose_list(ty));
422-
if x.is_empty() {
423-
let elem_ty = elt_hint.map_or_else(
424-
|| {
425-
if !self.solver().infer_with_first_use {
426-
self.error(
427-
errors,
428-
x.range(),
429-
ErrorInfo::Kind(ErrorKind::ImplicitAny),
430-
"This expression is implicitly inferred to be `list[Any]`. Please provide an explicit type annotation.".to_owned(),
431-
);
432-
Type::any_implicit()
433-
} else {
434-
self.solver().fresh_contained(self.uniques).to_type()
435-
}
436-
},
437-
|hint| hint.to_type(),
438-
);
439-
self.stdlib.list(elem_ty).to_type()
440-
} else {
441-
let elem_tys = self.elts_infer(&x.elts, elt_hint, errors);
442-
self.stdlib.list(self.unions(elem_tys)).to_type()
422+
if let Some(hint_ref) = hint.as_ref()
423+
&& let Type::Union(options) = hint_ref.ty()
424+
{
425+
for option in options {
426+
let branch_hint =
427+
self.decompose_list(HintRef::new(option, hint_ref.errors()));
428+
let ty = self.list_with_hint(x, branch_hint, errors);
429+
if self.is_subset_eq(&ty, option) {
430+
return ty;
431+
}
432+
}
443433
}
434+
let elt_hint = hint.and_then(|ty| self.decompose_list(ty));
435+
self.list_with_hint(x, elt_hint, errors)
444436
}
445437
Expr::Dict(x) => self.dict_infer(&x.items, hint, x.range, errors),
446438
Expr::Set(x) => {
@@ -1751,6 +1743,36 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
17511743
})
17521744
}
17531745

1746+
fn list_with_hint(
1747+
&self,
1748+
x: &ExprList,
1749+
elt_hint: Option<Hint>,
1750+
errors: &ErrorCollector,
1751+
) -> Type {
1752+
if x.is_empty() {
1753+
let elem_ty = elt_hint.map_or_else(
1754+
|| {
1755+
if !self.solver().infer_with_first_use {
1756+
self.error(
1757+
errors,
1758+
x.range(),
1759+
ErrorInfo::Kind(ErrorKind::ImplicitAny),
1760+
"This expression is implicitly inferred to be `list[Any]`. Please provide an explicit type annotation.".to_owned(),
1761+
);
1762+
Type::any_implicit()
1763+
} else {
1764+
self.solver().fresh_contained(self.uniques).to_type()
1765+
}
1766+
},
1767+
|hint| hint.to_type(),
1768+
);
1769+
self.stdlib.list(elem_ty).to_type()
1770+
} else {
1771+
let elem_tys = self.elts_infer(&x.elts, elt_hint, errors);
1772+
self.stdlib.list(self.unions(elem_tys)).to_type()
1773+
}
1774+
}
1775+
17541776
fn intercept_typing_self_use(&self, x: &Expr) -> Option<TypeInfo> {
17551777
match x {
17561778
Expr::Name(..) | Expr::Attribute(..) => {

‎pyrefly/lib/alt/unwrap.rs‎

Lines changed: 95 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,54 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
114114
}
115115
}
116116

117+
fn collect_var_from_hint<F>(&self, ty: &Type, make: &F) -> Option<Vec<Type>>
118+
where
119+
F: Fn(Var) -> Type,
120+
{
121+
match ty {
122+
Type::Union(tys) => {
123+
let mut collected = Vec::new();
124+
let mut matched = false;
125+
for branch in tys {
126+
if let Some(mut branch_res) = self.collect_var_from_hint(branch, make) {
127+
matched = true;
128+
collected.append(&mut branch_res);
129+
}
130+
}
131+
if matched { Some(collected) } else { None }
132+
}
133+
_ => {
134+
let var = self.fresh_var();
135+
let target = make(var);
136+
if self.is_subset_eq(&target, ty) {
137+
match self.resolve_var_opt(ty, var) {
138+
Some(resolved) => Some(vec![resolved]),
139+
None => Some(Vec::new()),
140+
}
141+
} else {
142+
None
143+
}
144+
}
145+
}
146+
}
147+
148+
fn hint_from_types<'b>(
149+
&self,
150+
mut types: Vec<Type>,
151+
hint: &HintRef<'b, '_>,
152+
) -> Option<Hint<'b>> {
153+
if types.is_empty() {
154+
None
155+
} else {
156+
let ty = if types.len() == 1 {
157+
types.pop().unwrap()
158+
} else {
159+
self.unions(types)
160+
};
161+
Some(hint.map_ty(|_| ty))
162+
}
163+
}
164+
117165
pub fn unwrap_mapping(&self, ty: &Type) -> Option<(Type, Type)> {
118166
let key = self.fresh_var();
119167
let value = self.fresh_var();
@@ -219,45 +267,65 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
219267
&self,
220268
hint: HintRef<'b, '_>,
221269
) -> (Option<Hint<'b>>, Option<Hint<'b>>) {
222-
let key = self.fresh_var();
223-
let value = self.fresh_var();
224-
let dict_type = self.stdlib.dict(key.to_type(), value.to_type()).to_type();
225-
if self.is_subset_eq(&dict_type, hint.ty()) {
226-
let key = hint.map_ty_opt(|ty| self.resolve_var_opt(ty, key));
227-
let value = hint.map_ty_opt(|ty| self.resolve_var_opt(ty, value));
228-
(key, value)
229-
} else {
230-
(None, None)
270+
let mut key_types = Vec::new();
271+
let mut value_types = Vec::new();
272+
let mut matched = false;
273+
274+
// Helper to process a single target type and accumulate results.
275+
let mut consider = |ty: &Type| {
276+
let key = self.fresh_var();
277+
let value = self.fresh_var();
278+
let dict_type = self.stdlib.dict(key.to_type(), value.to_type()).to_type();
279+
if self.is_subset_eq(&dict_type, ty) {
280+
matched = true;
281+
if let Some(key_ty) = self.resolve_var_opt(ty, key) {
282+
key_types.push(key_ty);
283+
}
284+
if let Some(value_ty) = self.resolve_var_opt(ty, value) {
285+
value_types.push(value_ty);
286+
}
287+
}
288+
};
289+
290+
match hint.ty() {
291+
Type::Union(branches) => {
292+
for branch in branches {
293+
consider(branch);
294+
}
295+
}
296+
ty => consider(ty),
231297
}
298+
299+
if !matched {
300+
return (None, None);
301+
}
302+
303+
let key = self.hint_from_types(key_types, &hint);
304+
let value = self.hint_from_types(value_types, &hint);
305+
(key, value)
232306
}
233307

234308
pub fn decompose_set<'b>(&self, hint: HintRef<'b, '_>) -> Option<Hint<'b>> {
235-
let elem = self.fresh_var();
236-
let set_type = self.stdlib.set(elem.to_type()).to_type();
237-
if self.is_subset_eq(&set_type, hint.ty()) {
238-
hint.map_ty_opt(|ty| self.resolve_var_opt(ty, elem))
239-
} else {
240-
None
309+
let make = |var: Var| self.stdlib.set(var.to_type()).to_type();
310+
match self.collect_var_from_hint(hint.ty(), &make) {
311+
Some(tys) => self.hint_from_types(tys, &hint),
312+
None => None,
241313
}
242314
}
243315

244316
pub fn decompose_list<'b>(&self, hint: HintRef<'b, '_>) -> Option<Hint<'b>> {
245-
let elem = self.fresh_var();
246-
let list_type = self.stdlib.list(elem.to_type()).to_type();
247-
if self.is_subset_eq(&list_type, hint.ty()) {
248-
hint.map_ty_opt(|ty| self.resolve_var_opt(ty, elem))
249-
} else {
250-
None
317+
let make = |var: Var| self.stdlib.list(var.to_type()).to_type();
318+
match self.collect_var_from_hint(hint.ty(), &make) {
319+
Some(tys) => self.hint_from_types(tys, &hint),
320+
None => None,
251321
}
252322
}
253323

254324
pub fn decompose_tuple<'b>(&self, hint: HintRef<'b, '_>) -> Option<Hint<'b>> {
255-
let elem = self.fresh_var();
256-
let tuple_type = self.stdlib.tuple(elem.to_type()).to_type();
257-
if self.is_subset_eq(&tuple_type, hint.ty()) {
258-
hint.map_ty_opt(|ty| self.resolve_var_opt(ty, elem))
259-
} else {
260-
None
325+
let make = |var: Var| self.stdlib.tuple(var.to_type()).to_type();
326+
match self.collect_var_from_hint(hint.ty(), &make) {
327+
Some(tys) => self.hint_from_types(tys, &hint),
328+
None => None,
261329
}
262330
}
263331

‎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)