Skip to content

Fix #1855: Multiassign from Union (take 2) #2219

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 30 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
c53a1c4
NewType+namedtuple: common method for typeinfo
elazarg Sep 3, 2016
7e043d0
handle union, add test
elazarg Sep 18, 2016
9a50d73
kill blank
elazarg Sep 18, 2016
5c4d86e
more tests
elazarg Sep 18, 2016
60cfbbb
handle binding
elazarg Sep 19, 2016
5e9e0f2
try to minimize visual difference
elazarg Sep 19, 2016
71f8475
(cont.)
elazarg Sep 19, 2016
42b6e73
add tests
elazarg Sep 21, 2016
0560bd8
no binder yet
elazarg Sep 23, 2016
da3a516
Support rebinding on multiassignment from union
elazarg Sep 27, 2016
ab60317
Merge remote-tracking branch 'upstream/master' into multiassign_union
elazarg Sep 27, 2016
6bb5519
more tests
elazarg Sep 27, 2016
3b4cd13
Merge remote-tracking branch 'upstream/master'
elazarg Oct 2, 2016
38651c4
handle union, add test
elazarg Sep 18, 2016
9000099
kill blank
elazarg Sep 18, 2016
f20f3d6
more tests
elazarg Sep 18, 2016
61be4e9
handle binding
elazarg Sep 19, 2016
9830cb4
try to minimize visual difference
elazarg Sep 19, 2016
59dc8b7
(cont.)
elazarg Sep 19, 2016
7f304e4
add tests
elazarg Sep 21, 2016
ff1ca80
no binder yet
elazarg Sep 23, 2016
168087e
Support rebinding on multiassignment from union
elazarg Sep 27, 2016
3f198cf
more tests
elazarg Sep 27, 2016
4fa059f
Rebase
elazarg Oct 5, 2016
ab35a4c
Merge
elazarg Oct 5, 2016
00f34ce
small merge fix
elazarg Oct 5, 2016
0cccb4b
Merge
elazarg Oct 7, 2016
a3ef3d5
what typeshed?
elazarg Oct 7, 2016
60c032f
Merge remote-tracking branch 'upstream/master' into multiassign_union2
elazarg Oct 11, 2016
21ac123
Merge remote-tracking branch 'upstream/master' into multiassign_union2
elazarg Oct 18, 2016
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion mypy/binder.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from typing import (Dict, List, Set, Iterator, Union)
from typing import (Dict, List, Set, Iterator, Optional, DefaultDict, Tuple, Union)
from contextlib import contextmanager
from collections import defaultdict

from mypy.types import Type, AnyType, PartialType
from mypy.nodes import (Key, Node, Expression, Var, RefExpr, SymbolTableNode)
Expand Down Expand Up @@ -45,6 +46,7 @@ class A:
reveal_type(lst[0].a) # str
```
"""
type_assignments = None # type: Optional[DefaultDict[Expression, List[Tuple[Type, Type]]]]

def __init__(self) -> None:
# The stack of frames currently used. These map
Expand Down Expand Up @@ -198,10 +200,20 @@ def get_declaration(self, expr: Node) -> Type:
else:
return None

@contextmanager
def accumulate_type_assignments(self) -> Iterator[DefaultDict[Expression,
List[Tuple[Type, Type]]]]:
self.type_assignments = defaultdict(list)
yield self.type_assignments
self.type_assignments = None

def assign_type(self, expr: Expression,
type: Type,
declared_type: Type,
restrict_any: bool = False) -> None:
if self.type_assignments is not None:
self.type_assignments[expr].append((type, declared_type))
return
if not expr.literal:
return
self.invalidate_dependencies(expr)
Expand Down
182 changes: 104 additions & 78 deletions mypy/checker.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@
from mypy.erasetype import erase_typevars
from mypy.expandtype import expand_type
from mypy.visitor import NodeVisitor
from mypy.join import join_types
from mypy.join import join_types, join_type_list
from mypy.treetransform import TransformVisitor
from mypy.meet import meet_simple, nearest_builtin_ancestor, is_overlapping_types
from mypy.binder import ConditionalTypeBinder
Expand Down Expand Up @@ -1070,8 +1070,15 @@ def check_assignment(self, lvalue: Lvalue, rvalue: Expression, infer_lvalue_type
new_syntax: bool = False) -> None:
"""Type check a single assignment: lvalue = rvalue."""
if isinstance(lvalue, TupleExpr) or isinstance(lvalue, ListExpr):
self.check_assignment_to_multiple_lvalues(lvalue.items, rvalue, lvalue,
infer_lvalue_type)
if isinstance(rvalue, TupleExpr) or isinstance(rvalue, ListExpr):
self.check_multi_assign_literal(lvalue.items, rvalue, lvalue, infer_lvalue_type)
return
# Infer the type of an ordinary rvalue expression.
# TODO maybe elsewhere; redundant
rvalue_type = self.accept(rvalue)
self.check_multi_assign(lvalue.items, rvalue, rvalue_type, lvalue,
undefined_rvalue=False,
infer_lvalue_type=infer_lvalue_type)
else:
lvalue_type, index_lvalue, inferred = self.check_lvalue(lvalue)
if lvalue_type:
Expand Down Expand Up @@ -1123,40 +1130,31 @@ def check_assignment(self, lvalue: Lvalue, rvalue: Expression, infer_lvalue_type
self.infer_variable_type(inferred, lvalue, self.accept(rvalue),
rvalue)

def check_assignment_to_multiple_lvalues(self, lvalues: List[Lvalue], rvalue: Expression,
context: Context,
infer_lvalue_type: bool = True) -> None:
if isinstance(rvalue, TupleExpr) or isinstance(rvalue, ListExpr):
# Recursively go into Tuple or List expression rhs instead of
# using the type of rhs, because this allowed more fine grained
# control in cases like: a, b = [int, str] where rhs would get
# type List[object]

rvalues = rvalue.items

if self.check_rvalue_count_in_assignment(lvalues, len(rvalues), context):
star_index = next((i for i, lv in enumerate(lvalues) if
isinstance(lv, StarExpr)), len(lvalues))

left_lvs = lvalues[:star_index]
star_lv = cast(StarExpr,
lvalues[star_index]) if star_index != len(lvalues) else None
right_lvs = lvalues[star_index + 1:]

left_rvs, star_rvs, right_rvs = self.split_around_star(
rvalues, star_index, len(lvalues))

lr_pairs = list(zip(left_lvs, left_rvs))
if star_lv:
rv_list = ListExpr(star_rvs)
rv_list.set_line(rvalue.get_line())
lr_pairs.append((star_lv.expr, rv_list))
lr_pairs.extend(zip(right_lvs, right_rvs))

for lv, rv in lr_pairs:
self.check_assignment(lv, rv, infer_lvalue_type)
else:
self.check_multi_assignment(lvalues, rvalue, context, infer_lvalue_type)
def check_multi_assign_literal(self, lvalues: List[Lvalue],
rvalue: Union[ListExpr, TupleExpr],
context: Context, infer_lvalue_type: bool = True) -> None:
# Recursively go into Tuple or List expression rhs instead of
# using the type of rhs, because this allowed more fine grained
# control in cases like: a, b = [int, str] where rhs would get
# type List[object]
# Tuple is also special cased to handle mutually nested lists and tuples
rvalues = rvalue.items
if self.check_rvalue_count_in_assignment(lvalues, len(rvalues), context):
star_index = next((i for (i, lv) in enumerate(lvalues) if isinstance(lv, StarExpr)),
len(lvalues))
left_lvs = lvalues[:star_index]
star_lv = cast(StarExpr, lvalues[star_index]) if star_index != len(lvalues) else None
right_lvs = lvalues[star_index + 1:]
left_rvs, star_rvs, right_rvs = self.split_around_star(
rvalues, star_index, len(lvalues))
lr_pairs = list(zip(left_lvs, left_rvs))
if star_lv:
rv_list = ListExpr(star_rvs)
rv_list.set_line(rvalue.get_line())
lr_pairs.append((star_lv.expr, rv_list))
lr_pairs.extend(zip(right_lvs, right_rvs))
for lv, rv in lr_pairs:
self.check_assignment(lv, rv, infer_lvalue_type)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nothing new in this function. It's only extracted, and the comment is slightly more elaborated.


def check_rvalue_count_in_assignment(self, lvalues: List[Lvalue], rvalue_count: int,
context: Context) -> bool:
Expand All @@ -1171,33 +1169,64 @@ def check_rvalue_count_in_assignment(self, lvalues: List[Lvalue], rvalue_count:
return False
return True

def check_multi_assignment(self, lvalues: List[Lvalue],
rvalue: Expression,
context: Context,
infer_lvalue_type: bool = True,
msg: str = None) -> None:
"""Check the assignment of one rvalue to a number of lvalues."""

# Infer the type of an ordinary rvalue expression.
rvalue_type = self.accept(rvalue) # TODO maybe elsewhere; redundant
undefined_rvalue = False

def check_multi_assign_from_any(self, lvalues: List[Expression], rvalue: Expression,
rvalue_type: AnyType, context: Context,
undefined_rvalue: bool,
infer_lvalue_type: bool) -> None:
for lv in lvalues:
if isinstance(lv, StarExpr):
lv = lv.expr
self.check_assignment(lv, self.temp_node(AnyType(), context), infer_lvalue_type)

def check_multi_assign(self, lvalues: List[Lvalue], rvalue: Expression,
rvalue_type: Type, context: Context, *,
undefined_rvalue: bool = False,
infer_lvalue_type: bool = True) -> None:
if isinstance(rvalue_type, AnyType):
for lv in lvalues:
if isinstance(lv, StarExpr):
lv = lv.expr
self.check_assignment(lv, self.temp_node(AnyType(), context), infer_lvalue_type)
self.check_multi_assign_from_any(lvalues, rvalue, rvalue_type,
context, undefined_rvalue, infer_lvalue_type)
elif isinstance(rvalue_type, TupleType):
self.check_multi_assignment_from_tuple(lvalues, rvalue, rvalue_type,
self.check_multi_assign_from_tuple(lvalues, rvalue, rvalue_type,
context, undefined_rvalue, infer_lvalue_type)
elif isinstance(rvalue_type, UnionType):
self.check_multi_assign_from_union(lvalues, rvalue, rvalue_type,
context, undefined_rvalue, infer_lvalue_type)
elif isinstance(rvalue_type, Instance) and self.instance_is_iterable(rvalue_type):
self.check_multi_assign_from_iterable(lvalues, rvalue, rvalue_type,
context, undefined_rvalue, infer_lvalue_type)
else:
self.check_multi_assignment_from_iterable(lvalues, rvalue_type,
context, infer_lvalue_type)
self.msg.type_not_iterable(rvalue_type, context)

def check_multi_assign_from_union(self, lvalues: List[Expression], rvalue: Expression,
rvalue_type: UnionType, context: Context,
undefined_rvalue: bool,
infer_lvalue_type: bool) -> None:
transposed = tuple([] for _ in lvalues) # type: Tuple[List[Type], ...]
with self.binder.accumulate_type_assignments() as assignments:
for item in rvalue_type.items:
self.check_multi_assign(lvalues, rvalue, item, context,
undefined_rvalue=True,
infer_lvalue_type=infer_lvalue_type)
for t, lv in zip(transposed, lvalues):
t.append(self.type_map.get(lv, AnyType()))
union_types = tuple(join_type_list(col) for col in transposed)
for expr, items in assignments.items():
types, declared_types = zip(*items)
self.binder.assign_type(expr,
join_type_list(types),
join_type_list(declared_types),
False)
for union, lv in zip(union_types, lvalues):
_1, _2, inferred = self.check_lvalue(lv)
if inferred:
self.set_inferred_type(inferred, lv, union)
else:
self.store_type(lv, union)

def check_multi_assignment_from_tuple(self, lvalues: List[Lvalue], rvalue: Expression,
rvalue_type: TupleType, context: Context,
undefined_rvalue: bool,
infer_lvalue_type: bool = True) -> None:
def check_multi_assign_from_tuple(self, lvalues: List[Lvalue], rvalue: Expression,
rvalue_type: TupleType, context: Context,
undefined_rvalue: bool,
infer_lvalue_type: bool) -> None:
if self.check_rvalue_count_in_assignment(lvalues, len(rvalue_type.items), context):
star_index = next((i for i, lv in enumerate(lvalues)
if isinstance(lv, StarExpr)), len(lvalues))
Expand Down Expand Up @@ -1275,25 +1304,22 @@ def split_around_star(self, items: List[T], star_index: int,
right = items[right_index:]
return (left, star, right)

def type_is_iterable(self, type: Type) -> bool:
return (is_subtype(type, self.named_generic_type('typing.Iterable',
[AnyType()])) and
isinstance(type, Instance))

def check_multi_assignment_from_iterable(self, lvalues: List[Lvalue], rvalue_type: Type,
context: Context,
infer_lvalue_type: bool = True) -> None:
if self.type_is_iterable(rvalue_type):
item_type = self.iterable_item_type(cast(Instance, rvalue_type))
for lv in lvalues:
if isinstance(lv, StarExpr):
self.check_assignment(lv.expr, self.temp_node(rvalue_type, context),
infer_lvalue_type)
else:
self.check_assignment(lv, self.temp_node(item_type, context),
infer_lvalue_type)
else:
self.msg.type_not_iterable(rvalue_type, context)
def instance_is_iterable(self, instance: Instance) -> bool:
return is_subtype(instance, self.named_generic_type('typing.Iterable',
[AnyType()]))

def check_multi_assign_from_iterable(self, lvalues: List[Expression], rvalue: Expression,
rvalue_type: Instance, context: Context,
undefined_rvalue: bool,
infer_lvalue_type: bool) -> None:
item_type = self.iterable_item_type(rvalue_type)
for lv in lvalues:
if isinstance(lv, StarExpr):
self.check_assignment(lv.expr, self.temp_node(rvalue_type, context),
infer_lvalue_type)
else:
self.check_assignment(lv, self.temp_node(item_type, context),
infer_lvalue_type)

def check_lvalue(self, lvalue: Lvalue) -> Tuple[Type, IndexExpr, Var]:
lvalue_type = None # type: Type
Expand Down
120 changes: 120 additions & 0 deletions test-data/unit/check-unions.test
Original file line number Diff line number Diff line change
Expand Up @@ -129,3 +129,123 @@ class C(Generic[T, U]):
a = C() # type: C[int, int]
b = a.f('a')
a.f(b) # E: Argument 1 to "f" of "C" has incompatible type "int"; expected "str"

[case testUnionMultiassign1]
from typing import Union, Tuple, Any

b = None # type: Union[Tuple[int], Tuple[float]]
(b1,) = b
reveal_type(b1) # E: Revealed type is 'builtins.float'

[case testUnionMultiassign2]
from typing import Union, Tuple

c = None # type: Union[Tuple[int, int], Tuple[int, str]]
q = None # type: Tuple[int, float]
(c1, c2) = c
reveal_type(c1) # E: Revealed type is 'builtins.int'
reveal_type(c2) # E: Revealed type is 'builtins.object'

[case testUnionMultiassignAny]
from typing import Union, Tuple, Any

d = None # type: Union[Any, Tuple[float, float]]
(d1, d2) = d
reveal_type(d1) # E: Revealed type is 'Any'
reveal_type(d2) # E: Revealed type is 'Any'

e = None # type: Union[Any, Tuple[float, float], int]
(e1, e2) = e # E: 'builtins.int' object is not iterable

[case testUnionMultiassignJoin]
from typing import Union, List

class A: pass
class B(A): pass
class C(A): pass
a = None # type: Union[List[B], List[C]]
x, y = a
reveal_type(x) # E: Revealed type is '__main__.A'

[builtins fixtures/list.pyi]
[case testUnionMultiassignRebind]
from typing import Union, List

class A: pass
class B(A): pass
class C(A): pass
obj = None # type: object
c = None # type: object
a = None # type: Union[List[B], List[C]]
obj, new = a
reveal_type(obj) # E: Revealed type is '__main__.A'
reveal_type(new) # E: Revealed type is '__main__.A'

obj = 1
reveal_type(obj) # E: Revealed type is 'builtins.int'

[builtins fixtures/list.pyi]

[case testUnionMultiassignAlreadyDeclared]
from typing import Union, Tuple

a = None # type: Union[Tuple[int, int], Tuple[int, float]]
a1 = None # type: object
a2 = None # type: int

(a1, a2) = a # E: Incompatible types in assignment (expression has type "float", variable has type "int")

b = None # type: Union[Tuple[float, int], Tuple[int, int]]
b1 = None # type: object
b2 = None # type: int

(b1, b2) = a # E: Incompatible types in assignment (expression has type "float", variable has type "int")

c = None # type: Union[Tuple[int, int], Tuple[int, int]]
c1 = None # type: object
c2 = None # type: int

(c1, c2) = c
reveal_type(c1) # E: Revealed type is 'builtins.int'
reveal_type(c2) # E: Revealed type is 'builtins.int'

d = None # type: Union[Tuple[int, int], Tuple[int, float]]
d1 = None # type: object

(d1, d2) = d
reveal_type(d1) # E: Revealed type is 'builtins.int'
reveal_type(d2) # E: Revealed type is 'builtins.float'

[case testUnionMultiassignIndexed]
from typing import Union, Tuple, List

class B:
x = None # type: object

x = None # type: List[int]
b = None # type: B

a = None # type: Union[Tuple[int, int], Tuple[int, object]]
(x[0], b.x) = a

# I don't know why is it incomplete type
reveal_type(x[0]) # E: Revealed type is 'builtins.int*'
reveal_type(b.x) # E: Revealed type is 'builtins.object'

[builtins fixtures/list.pyi]

[case testUnionMultiassignPacked]
from typing import Union, Tuple, List

a = None # type: Union[Tuple[int, int, int], Tuple[int, int, str]]
a1 = None # type: int
a2 = None # type: object
--FIX: allow proper rebinding of packed
xs = None # type: List[int]
(a1, *xs, a2) = a

reveal_type(a1) # E: Revealed type is 'builtins.int'
reveal_type(xs) # E: Revealed type is 'builtins.list[builtins.int]'
reveal_type(a2) # E: Revealed type is 'builtins.int'

[builtins fixtures/list.pyi]
1 change: 1 addition & 0 deletions test-data/unit/fixtures/list.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ class list(Iterable[T], Generic[T]):
def __add__(self, x: list[T]) -> list[T]: pass
def __mul__(self, x: int) -> list[T]: pass
def __getitem__(self, x: int) -> T: pass
def __setitem__(self, x: int, v: T) -> None: pass
def append(self, x: T) -> None: pass
def extend(self, x: Iterable[T]) -> None: pass

Expand Down