Skip to content

Commit 7fa4425

Browse files
authored
Improve type inference by supporting partial OrderedDict types (#8039)
This example no longer needs a type annotation: ``` from collections import OrderedDict x = OrderedDict() x[1] = 'x' ``` Work towards #1055.
1 parent f01eccd commit 7fa4425

File tree

4 files changed

+25
-4
lines changed

4 files changed

+25
-4
lines changed

mypy/checker.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2809,7 +2809,8 @@ def infer_partial_type(self, name: Var, lvalue: Lvalue, init_type: Type) -> bool
28092809
if (isinstance(lvalue, (NameExpr, MemberExpr)) and
28102810
(fullname == 'builtins.list' or
28112811
fullname == 'builtins.set' or
2812-
fullname == 'builtins.dict') and
2812+
fullname == 'builtins.dict' or
2813+
fullname == 'collections.OrderedDict') and
28132814
all(isinstance(t, (NoneType, UninhabitedType))
28142815
for t in get_proper_types(init_type.args))):
28152816
partial_type = PartialType(init_type.type, name, init_type.args)
@@ -3002,7 +3003,7 @@ def try_infer_partial_type_from_indexed_assignment(
30023003
if partial_types is None:
30033004
return
30043005
typename = type_type.fullname
3005-
if typename == 'builtins.dict':
3006+
if typename == 'builtins.dict' or typename == 'collections.OrderedDict':
30063007
# TODO: Don't infer things twice.
30073008
key_type = self.expr_checker.accept(lvalue.index)
30083009
value_type = self.expr_checker.accept(rvalue)
@@ -3013,7 +3014,7 @@ def try_infer_partial_type_from_indexed_assignment(
30133014
if (is_valid_inferred_type(full_key_type) and
30143015
is_valid_inferred_type(full_value_type)):
30153016
if not self.current_node_deferred:
3016-
var.type = self.named_generic_type('builtins.dict',
3017+
var.type = self.named_generic_type(typename,
30173018
[full_key_type, full_value_type])
30183019
del partial_types[var]
30193020

mypy/checkexpr.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -543,6 +543,7 @@ def check_typeddict_call_with_kwargs(self, callee: TypedDictType,
543543
} # type: ClassVar[Dict[str, List[str]]]
544544
container_args = {'builtins.list': {'extend': ['builtins.list']},
545545
'builtins.dict': {'update': ['builtins.dict']},
546+
'collections.OrderedDict': {'update': ['builtins.dict']},
546547
'builtins.set': {'update': ['builtins.set', 'builtins.list']},
547548
} # type: ClassVar[Dict[str, Dict[str, List[str]]]]
548549

test-data/unit/check-inference.test

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1572,6 +1572,19 @@ if bool():
15721572
reveal_type(dd) # N: Revealed type is 'builtins.dict[Any, Any]'
15731573
[builtins fixtures/dict.pyi]
15741574

1575+
[case testInferOrderedDictInitializedToEmpty]
1576+
from collections import OrderedDict
1577+
1578+
o = OrderedDict()
1579+
o[1] = 'x'
1580+
reveal_type(o) # N: Revealed type is 'collections.OrderedDict[builtins.int, builtins.str]'
1581+
1582+
d = {1: 'x'}
1583+
oo = OrderedDict()
1584+
oo.update(d)
1585+
reveal_type(oo) # N: Revealed type is 'collections.OrderedDict[builtins.int*, builtins.str*]'
1586+
[builtins fixtures/dict.pyi]
1587+
15751588

15761589
-- Inferring types of variables first initialized to None (partial types)
15771590
-- ----------------------------------------------------------------------

test-data/unit/lib-stub/collections.pyi

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import Any, Iterable, Union, Optional
1+
from typing import Any, Iterable, Union, Optional, Dict, TypeVar
22

33
def namedtuple(
44
typename: str,
@@ -9,3 +9,9 @@ def namedtuple(
99
module: Optional[str] = ...,
1010
defaults: Optional[Iterable[Any]] = ...
1111
) -> Any: ...
12+
13+
K = TypeVar('K')
14+
V = TypeVar('V')
15+
16+
class OrderedDict(Dict[K, V]):
17+
def __setitem__(self, k: K, v: V) -> None: ...

0 commit comments

Comments
 (0)