Improve type inference by supporting partial OrderedDict types

mypy/checker.py

@@ -2809,7 +2809,8 @@ def infer_partial_type(self, name: Var, lvalue: Lvalue, init_type: Type) -> bool
             if (isinstance(lvalue, (NameExpr, MemberExpr)) and
                     (fullname == 'builtins.list' or
                      fullname == 'builtins.set' or
-                     fullname == 'builtins.dict') and
+                     fullname == 'builtins.dict' or
+                     fullname == 'collections.OrderedDict') and
                     all(isinstance(t, (NoneType, UninhabitedType))
                         for t in get_proper_types(init_type.args))):
                 partial_type = PartialType(init_type.type, name, init_type.args)
@@ -3002,7 +3003,7 @@ def try_infer_partial_type_from_indexed_assignment(
                 if partial_types is None:
                     return
                 typename = type_type.fullname
-                if typename == 'builtins.dict':
+                if typename == 'builtins.dict' or typename == 'collections.OrderedDict':
                     # TODO: Don't infer things twice.
                     key_type = self.expr_checker.accept(lvalue.index)
                     value_type = self.expr_checker.accept(rvalue)
@@ -3013,7 +3014,7 @@ def try_infer_partial_type_from_indexed_assignment(
                     if (is_valid_inferred_type(full_key_type) and
                             is_valid_inferred_type(full_value_type)):
                         if not self.current_node_deferred:
-                            var.type = self.named_generic_type('builtins.dict',
+                            var.type = self.named_generic_type(typename,
                                                                [full_key_type, full_value_type])
                             del partial_types[var]
 

mypy/checkexpr.py

@@ -543,6 +543,7 @@ def check_typeddict_call_with_kwargs(self, callee: TypedDictType,
                  }  # type: ClassVar[Dict[str, List[str]]]
     container_args = {'builtins.list': {'extend': ['builtins.list']},
                       'builtins.dict': {'update': ['builtins.dict']},
+                      'collections.OrderedDict': {'update': ['builtins.dict']},
                       'builtins.set': {'update': ['builtins.set', 'builtins.list']},
                       }  # type: ClassVar[Dict[str, Dict[str, List[str]]]]
 

test-data/unit/check-inference.test

@@ -1572,6 +1572,19 @@ if bool():
 reveal_type(dd) # N: Revealed type is 'builtins.dict[Any, Any]'
 [builtins fixtures/dict.pyi]
 
+[case testInferOrderedDictInitializedToEmpty]
+from collections import OrderedDict
+
+o = OrderedDict()
+o[1] = 'x'
+reveal_type(o) # N: Revealed type is 'collections.OrderedDict[builtins.int, builtins.str]'
+
+d = {1: 'x'}
+oo = OrderedDict()
+oo.update(d)
+reveal_type(oo) # N: Revealed type is 'collections.OrderedDict[builtins.int*, builtins.str*]'
+[builtins fixtures/dict.pyi]
+
 
 -- Inferring types of variables first initialized to None (partial types)
 -- ----------------------------------------------------------------------

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

@@ -1,4 +1,4 @@
-from typing import Any, Iterable, Union, Optional
+from typing import Any, Iterable, Union, Optional, Dict, TypeVar
 
 def namedtuple(
     typename: str,
@@ -9,3 +9,9 @@ def namedtuple(
     module: Optional[str] = ...,
     defaults: Optional[Iterable[Any]] = ...
 ) -> Any: ...
+
+K = TypeVar('K')
+V = TypeVar('V')
+
+class OrderedDict(Dict[K, V]):
+    def __setitem__(self, k: K, v: V) -> None: ...