Repository: python/mypy

Issue:
After Optional[str] == 'literal', infer that it's not optional
Example:

``` python
def f(a: Optional[str]) -> str:
    if a == '<string>': return a[1:-1]
    else: return ''
```

With `--strict-optional`, this produces an error:

```
error: Value of type None is not indexable
```

It would be nice if mypy understood that a is a string if this == returns True.



Discussion:
We should probably fix this more generally with Unions:

``` py
def f(a: Union[int, str]) -> str
   if a == '<string>': return a  # mypy should understand that a must be a string here
   else: return str(a)
```

Are we okay with saying that equality must imply some sort of type equivalence?  (In particular, that `a == b` implies `a` is either a super- or sub-class of `b`.)

I was more thinking that the join of the two types should not be object.

We should carefully introduce and test such a change because I expect it'll cause some noise in various large projects (most of which I expect to be questionable code in those projects, but some might be due to our rules needing some exception, and either wait it'll take a while to sort it all out).

FWIW the same rule should apply to `x in c` and `x not in c` -- the type of x and the type of c's items must have a non-trivial join.

Asuming that on `a == 'text'`, a is string is somewhat unsafe. There's a slight danger given that == can be redefined, so I could have a class with an `__eq__` method that considers itself equal to string but doesn't have other string properties. A similar problem occurs with "in" given that it depends on `__eq__`.

The only operator where I think this could safely be done is "is" and "is not", but that doesn't cover the example.

In general that assumption is unsafe, but here `a` has type `Optional[str] = Union[None, str]` and it's quite reasonable to assume that `==` will not be redefined on `None`.


Resolving patch:
diff --git a/mypy/checker.py b/mypy/checker.py
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -2436,6 +2436,17 @@ def is_literal_none(n: Expression) -> bool:
     return isinstance(n, NameExpr) and n.fullname == 'builtins.None'
 
 
+def is_optional(t: Type) -> bool:
+    return isinstance(t, UnionType) and any(isinstance(e, NoneTyp) for e in t.items)
+
+
+def remove_optional(typ: Type) -> Type:
+    if isinstance(typ, UnionType):
+        return UnionType.make_union([t for t in typ.items if not isinstance(t, NoneTyp)])
+    else:
+        return typ
+
+
 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
@@ -2505,11 +2516,10 @@ def find_isinstance_check(node: Expression,
                 vartype = type_map[expr]
                 type = get_isinstance_type(node.args[1], type_map)
                 return conditional_type_map(expr, vartype, type)
-    elif (isinstance(node, ComparisonExpr) and any(is_literal_none(n) for n in node.operands) and
-          experiments.STRICT_OPTIONAL):
+    elif (isinstance(node, ComparisonExpr) and experiments.STRICT_OPTIONAL):
         # Check for `x is None` and `x is not None`.
         is_not = node.operators == ['is not']
-        if is_not or node.operators == ['is']:
+        if any(is_literal_none(n) for n in node.operands) and (is_not or node.operators == ['is']):
             if_vars = {}  # type: Dict[Expression, Type]
             else_vars = {}  # type: Dict[Expression, Type]
             for expr in node.operands:
@@ -2524,6 +2534,20 @@ def find_isinstance_check(node: Expression,
             if is_not:
                 if_vars, else_vars = else_vars, if_vars
             return if_vars, else_vars
+        # Check for `x == y` where x is of type Optional[T] and y is of type T
+        # or a type that overlaps with T (or vice versa).
+        elif node.operators == ['==']:
+            first_type = type_map[node.operands[0]]
+            second_type = type_map[node.operands[1]]
+            if is_optional(first_type) != is_optional(second_type):
+                if is_optional(first_type):
+                    optional_type, comp_type = first_type, second_type
+                    optional_expr = node.operands[0]
+                else:
+                    optional_type, comp_type = second_type, first_type
+                    optional_expr = node.operands[1]
+                if is_overlapping_types(optional_type, comp_type):
+                    return {optional_expr: remove_optional(optional_type)}, {}
     elif isinstance(node, RefExpr):
         # Restrict the type of the variable to True-ish/False-ish in the if and else branches
         # respectively


Test patch:
diff --git a/test-data/unit/check-optional.test b/test-data/unit/check-optional.test
--- a/test-data/unit/check-optional.test
+++ b/test-data/unit/check-optional.test
@@ -451,3 +451,49 @@ reveal_type(l)  # E: Revealed type is 'builtins.list[typing.Generator*[builtins.
 [case testNoneListTernary]
 x = [None] if "" else [1]  # E: List item 0 has incompatible type "int"
 [builtins fixtures/list.pyi]
+
+[case testInferEqualsNotOptional]
+from typing import Optional
+x = ''  # type: Optional[str]
+if x == '<string>':
+    reveal_type(x)  # E: Revealed type is 'builtins.str'
+else:
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.str, builtins.None]'
+[builtins fixtures/ops.pyi]
+
+[case testInferEqualsNotOptionalWithUnion]
+from typing import Union
+x = ''  # type: Union[str, int, None]
+if x == '<string>':
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.str, builtins.int]'
+else:
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.str, builtins.int, builtins.None]'
+[builtins fixtures/ops.pyi]
+
+[case testInferEqualsNotOptionalWithOverlap]
+from typing import Union
+x = ''  # type: Union[str, int, None]
+if x == object():
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.str, builtins.int]'
+else:
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.str, builtins.int, builtins.None]'
+[builtins fixtures/ops.pyi]
+
+[case testInferEqualsStillOptionalWithNoOverlap]
+from typing import Optional
+x = ''  # type: Optional[str]
+if x == 0:
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.str, builtins.None]'
+else:
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.str, builtins.None]'
+[builtins fixtures/ops.pyi]
+
+[case testInferEqualsStillOptionalWithBothOptional]
+from typing import Union
+x = ''  # type: Union[str, int, None]
+y = ''  # type: Union[str, None]
+if x == y:
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.str, builtins.int, builtins.None]'
+else:
+    reveal_type(x)  # E: Revealed type is 'Union[builtins.str, builtins.int, builtins.None]'
+[builtins fixtures/ops.pyi]