Allow optional removal after container 'in' checks

python · JukkaL · Oct 24, 2017 · Sep 14, 2017 · Sep 14, 2017 · Sep 14, 2017
commit 363407c6115fd0365da8d6e27f91d44a900c3d78
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -2702,12 +2702,10 @@ def iterable_item_type(self, instance: Instance) -> Type:
     def function_type(self, func: FuncBase) -> FunctionLike:
         return function_type(func, self.named_type('builtins.function'))
 
-    # TODO: These next two functions should refer to TypeMap below
-    def find_isinstance_check(self, n: Expression) -> Tuple[Optional[Dict[Expression, Type]],
-                                                            Optional[Dict[Expression, Type]]]:
+    def find_isinstance_check(self, n: Expression) -> 'Tuple[TypeMap, TypeMap]':
         return find_isinstance_check(n, self.type_map)
 
-    def push_type_map(self, type_map: Optional[Dict[Expression, Type]]) -> None:
+    def push_type_map(self, type_map: 'TypeMap') -> None:
         if type_map is None:
             self.binder.unreachable()
         else:
@@ -2874,6 +2872,16 @@ def remove_optional(typ: Type) -> Type:
         return typ
 
 
+def is_non_optional_container(tp: Type) -> bool:
+    # This is only safe for built-in containers, where we know the behavior of __contains__.
+    if not isinstance(tp, Instance):
+        return False
+    if tp.type.fullname() not in ['builtins.list', 'builtins.dict',
+                                  'builtins.set', 'builtins.frozenfet']:
+        return False
+    return not is_optional(tp.args[0])
+
+
 def and_conditional_maps(m1: TypeMap, m2: TypeMap) -> TypeMap:
     """Calculate what information we can learn from the truth of (e1 and e2)
     in terms of the information that we can learn from the truth of e1 and
@@ -3020,6 +3028,11 @@ def find_isinstance_check(node: Expression,
                     optional_expr = node.operands[1]
                 if is_overlapping_types(optional_type, comp_type):
                     return {optional_expr: remove_optional(optional_type)}, {}
+        elif node.operators == ['in']:
+            item_type = type_map[node.operands[0]]
+            cont_type = type_map[node.operands[1]]
+            if is_optional(item_type) and is_non_optional_container(cont_type):
+                return {node.operands[0]: remove_optional(item_type)}, {}
     elif isinstance(node, RefExpr):
         # Restrict the type of the variable to True-ish/False-ish in the if and else branches
         # respectively

diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -450,7 +450,7 @@ def prepare_method_signature(self, func: FuncDef) -> None:
                         leading_type = fill_typevars(self.type)
                     func.type = replace_implicit_first_type(functype, leading_type)
 
-    def set_original_def(self, previous: Node, new: FuncDef) -> bool:
+    def set_original_def(self, previous: Optional[Node], new: FuncDef) -> bool:
         """If 'new' conditionally redefine 'previous', set 'previous' as original
 
         We reject straight redefinitions of functions, as they are usually
@@ -2650,7 +2650,7 @@ def check_classvar(self, s: AssignmentStmt) -> None:
         lvalue = s.lvalues[0]
         if len(s.lvalues) != 1 or not isinstance(lvalue, RefExpr):
             return
-        if not self.is_classvar(s.type):
+        if not s.type or not self.is_classvar(s.type):
             return
         if self.is_class_scope() and isinstance(lvalue, NameExpr):
             node = lvalue.node
@@ -2661,7 +2661,7 @@ def check_classvar(self, s: AssignmentStmt) -> None:
             # Other kinds of member assignments should be already reported
             self.fail_invalid_classvar(lvalue)
 
-    def is_classvar(self, typ: Optional[Type]) -> bool:
+    def is_classvar(self, typ: Type) -> bool:
         if not isinstance(typ, UnboundType):
             return False
         sym = self.lookup_qualified(typ.name, typ)
@@ -3710,12 +3710,13 @@ def is_module_scope(self) -> bool:
     def add_symbol(self, name: str, node: SymbolTableNode,
                    context: Context) -> None:
         if self.is_func_scope():
-            if name in self.locals[-1]:
+            localns = self.locals[-1]
+            assert localns is not None
+            if name in localns:
                 # Flag redefinition unless this is a reimport of a module.
-                if not (node.kind == MODULE_REF and
-                        self.locals[-1][name].node == node.node):
+                if not (node.kind == MODULE_REF and localns[name].node == node.node):
                     self.name_already_defined(name, context)
-            self.locals[-1][name] = node
+            localns[name] = node
         elif self.type:
             self.type.names[name] = node
         else:
@@ -3733,11 +3734,13 @@ def add_symbol(self, name: str, node: SymbolTableNode,
             self.globals[name] = node
 
     def add_local(self, node: Union[Var, FuncDef, OverloadedFuncDef], ctx: Context) -> None:
+        localns = self.locals[-1]
+        assert localns is not None, "Should not add locals outside a function"
         name = node.name()
-        if name in self.locals[-1]:
+        if name in localns:
             self.name_already_defined(name, ctx)
         node._fullname = name
-        self.locals[-1][name] = SymbolTableNode(LDEF, node)
+        localns[name] = SymbolTableNode(LDEF, node)
 
     def add_exports(self, *exps: Expression) -> None:
         for exp in exps:
@@ -4203,7 +4206,7 @@ def visit_decorator(self, dec: Decorator) -> None:
         for expr in dec.decorators:
             preserve_type = False
             if isinstance(expr, RefExpr) and isinstance(expr.node, FuncDef):
-                if is_identity_signature(expr.node.type):
+                if expr.node.type and is_identity_signature(expr.node.type):
                     preserve_type = True
             if not preserve_type:
                 decorator_preserves_type = False
@@ -4297,14 +4300,17 @@ def perform_transform(self, node: Union[Node, SymbolTableNode],
                           transform: Callable[[Type], Type]) -> None:
         """Apply transform to all types associated with node."""
         if isinstance(node, ForStmt):
-            node.index_type = transform(node.index_type)
+            if node.index_type:
+                node.index_type = transform(node.index_type)
             self.transform_types_in_lvalue(node.index, transform)
         if isinstance(node, WithStmt):
-            node.target_type = transform(node.target_type)
+            if node.target_type:
+                node.target_type = transform(node.target_type)
             for n in node.target:
                 if isinstance(n, NameExpr) and isinstance(n.node, Var) and n.node.type:
                     n.node.type = transform(n.node.type)
         if isinstance(node, (FuncDef, CastExpr, AssignmentStmt, TypeAliasExpr, Var)):
+            assert node.type, "Scheduled patch for non-existent type"
             node.type = transform(node.type)
         if isinstance(node, NewTypeExpr):
             node.old_type = transform(node.old_type)
@@ -4314,14 +4320,17 @@ def perform_transform(self, node: Union[Node, SymbolTableNode],
             if node.values:
                 node.values = [transform(v) for v in node.values]
         if isinstance(node, TypedDictExpr):
+            assert node.info.typeddict_type, "Scheduled patch for non-existent type"
             node.info.typeddict_type = cast(TypedDictType,
                                             transform(node.info.typeddict_type))
         if isinstance(node, NamedTupleExpr):
+            assert node.info.tuple_type, "Scheduled patch for non-existent type"
             node.info.tuple_type = cast(TupleType,
                                         transform(node.info.tuple_type))
         if isinstance(node, TypeApplication):
             node.types = [transform(t) for t in node.types]
         if isinstance(node, SymbolTableNode):
+            assert node.type_override, "Scheduled patch for non-existent type"
             node.type_override = transform(node.type_override)
         if isinstance(node, TypeInfo):
             for tvar in node.defn.type_vars:
@@ -4347,7 +4356,8 @@ def transform_types_in_lvalue(self, lvalue: Lvalue,
         if isinstance(lvalue, RefExpr):
             if isinstance(lvalue.node, Var):
                 var = lvalue.node
-                var.type = transform(var.type)
+                if var.type:
+                    var.type = transform(var.type)
         elif isinstance(lvalue, TupleExpr):
             for item in lvalue.items:
                 self.transform_types_in_lvalue(item, transform)