Repository: python/mypy

Issue:
Type context not working for second argument of dict.get
The new signature for `dict.get` introduced in https://github.com/python/typeshed/commit/6008b9dbb14bd947e219f655c66a36a63a91c42e causes wrong types to be inferred for some arguments. Here is an example that demonstrates the issue and doesn't rely on the typeshed commit:

```
from typing import Any, Generic, overload, Optional, TypeVar, Union

_KT = TypeVar('_KT')
_VT_co = TypeVar('_VT_co', covariant=True)
_T = TypeVar('_T')

class D(Generic[_KT, _VT_co]):  # Mimics Dict
    @overload
    def get(self, k: _KT) -> Optional[_VT_co]: ...
    @overload
    def get(self, k: _KT, default: _T) -> Union[_VT_co, _T]: ...

def f(x: Any, y: D[int, D[int, Any]]) -> None:
    a = y.get(x, D())
    reveal_type(a)  # Union[t.D[builtins.int, Any], t.D*[builtins.None, builtins.None]]
```

The expected type is `t.D[builtins.int, Any]`.

Similar example using typeshed (note that the commit that introduced the problem is not yet used in mypy master):

```
from typing import Any, Dict

def f(x: Any, y: Dict[int, Dict[int, Any]]) -> None:
    a = y.get(x, {})
    reveal_type(a)  # Union[builtins.dict[builtins.int, Any], builtins.dict*[builtins.None, builtins.None]]
```

The reason for the behavior is that mypy doesn't realize that type context should propagate from the return type to the argument somehow. Actually, it's unclear whether type context should generally work like this -- perhaps we should modify the signature of `get` instead.


Discussion:
The first example doesn't even need the 1-arg overload. The covariance of the value type also seems unrelated.

But I'm not sure I understand what you think the behavior should be. The type of `default` is intentionally a different type variable than D's value type. So what is the type of `D()` passed as the default? Why should it conform to the value type?

I tried to rewrite the example using a separate assignment:
```python
    d = D()
    a = y.get(x, d)
```
but this says `Need type annotation for variable` on the `d = D()` line. But somehow when I change to
```python
    d = D()  # type: D
    a = y.get(x, d)
```
all is well. So it looks like there's something half-hearted being done to the `D()` call in your original version, where it is not quite given the type `D[Any, Any]`?

If you think there's simply a bug I'm happy to sit back and look at your fix to see if I can understand the issue that way. :-) (But your last sentence made it sound like you're not sure yourself?)
The explicit annotation works because `D` is a shorthand for `D[Any, Any]`, so the `None` types get never inferred.

There are at least two ways in which we could fix this:

1. Somehow propagate information back from the union return type to the type context, with the goal of simplifying away the union in the return type. The rationale here is that it's preferable to have a non-union return than a union return type if we are free to choose either by substituting type variables suitably.
2. Update the second argument to `get` to have type `Union[_VT_co, _T]` and pick up the `_VT_co` context from there. I checked that this doesn't work yet, so this will also require a mypy change.

The first option sounds pretty complicated, so I'm planning to try the second option first.
OK, now I see. It looks like the current state of `Mapping.get()` etc. causes lots of issues with real-world code, so I agree on the priority. Solution (2) sounds fine.

Resolving patch:
diff --git a/mypy/constraints.py b/mypy/constraints.py
--- a/mypy/constraints.py
+++ b/mypy/constraints.py
@@ -2,10 +2,12 @@
 
 from typing import Iterable, List, Optional
 
+from mypy import experiments
 from mypy.types import (
     CallableType, Type, TypeVisitor, UnboundType, AnyType, Void, NoneTyp, TypeVarType,
     Instance, TupleType, TypedDictType, UnionType, Overloaded, ErasedType, PartialType,
-    DeletedType, UninhabitedType, TypeType, TypeVarId, is_named_instance
+    DeletedType, UninhabitedType, TypeType, TypeVarId, TypeQuery, ALL_TYPES_STRATEGY,
+    is_named_instance
 )
 from mypy.maptype import map_instance_to_supertype
 from mypy import nodes
@@ -149,13 +151,23 @@ def infer_constraints(template: Type, actual: Type,
     # be a supertype of the potential subtype, some item of the Union
     # must be a supertype of it.
     if direction == SUBTYPE_OF and isinstance(actual, UnionType):
+        # If some of items is not a complete type, disregard that.
+        items = simplify_away_incomplete_types(actual.items)
+        # We infer constraints eagerly -- try to find constraints for a type
+        # variable if possible. This seems to help with some real-world
+        # use cases.
         return any_constraints(
             [infer_constraints_if_possible(template, a_item, direction)
-             for a_item in actual.items])
+             for a_item in items],
+            eager=True)
     if direction == SUPERTYPE_OF and isinstance(template, UnionType):
+        # When the template is a union, we are okay with leaving some
+        # type variables indeterminate. This helps with some special
+        # cases, though this isn't very principled.
         return any_constraints(
             [infer_constraints_if_possible(t_item, actual, direction)
-             for t_item in template.items])
+             for t_item in template.items],
+            eager=False)
 
     # Remaining cases are handled by ConstraintBuilderVisitor.
     return template.accept(ConstraintBuilderVisitor(actual, direction))
@@ -177,12 +189,18 @@ def infer_constraints_if_possible(template: Type, actual: Type,
     return infer_constraints(template, actual, direction)
 
 
-def any_constraints(options: List[Optional[List[Constraint]]]) -> List[Constraint]:
-    """Deduce what we can from a collection of constraint lists given that
-    at least one of the lists must be satisfied. A None element in the
-    list of options represents an unsatisfiable constraint and is ignored.
+def any_constraints(options: List[Optional[List[Constraint]]], eager: bool) -> List[Constraint]:
+    """Deduce what we can from a collection of constraint lists.
+
+    It's a given that at least one of the lists must be satisfied. A
+    None element in the list of options represents an unsatisfiable
+    constraint and is ignored.  Ignore empty constraint lists if eager
+    is true -- they are always trivially satisfiable.
     """
-    valid_options = [option for option in options if option is not None]
+    if eager:
+        valid_options = [option for option in options if option]
+    else:
+        valid_options = [option for option in options if option is not None]
     if len(valid_options) == 1:
         return valid_options[0]
     # Otherwise, there are either no valid options or multiple valid options.
@@ -196,6 +214,34 @@ def any_constraints(options: List[Optional[List[Constraint]]]) -> List[Constrain
     # every option, combine the bounds with meet/join.
 
 
+def simplify_away_incomplete_types(types: List[Type]) -> List[Type]:
+    complete = [typ for typ in types if is_complete_type(typ)]
+    if complete:
+        return complete
+    else:
+        return types
+
+
+def is_complete_type(typ: Type) -> bool:
+    """Is a type complete?
+
+    A complete doesn't have uninhabited type components or (when not in strict
+    optional mode) None components.
+    """
+    return typ.accept(CompleteTypeVisitor())
+
+
+class CompleteTypeVisitor(TypeQuery):
+    def __init__(self) -> None:
+        super().__init__(default=True, strategy=ALL_TYPES_STRATEGY)
+
+    def visit_none_type(self, t: NoneTyp) -> bool:
+        return experiments.STRICT_OPTIONAL
+
+    def visit_uninhabited_type(self, t: UninhabitedType) -> bool:
+        return False
+
+
 class ConstraintBuilderVisitor(TypeVisitor[List[Constraint]]):
     """Visitor class for inferring type constraints."""
 


Test patch:
diff --git a/test-data/unit/check-inference-context.test b/test-data/unit/check-inference-context.test
--- a/test-data/unit/check-inference-context.test
+++ b/test-data/unit/check-inference-context.test
@@ -826,3 +826,42 @@ class A(Generic[T]):
     pass
 reveal_type(A()) # E: Revealed type is '__main__.A[builtins.None]'
 b = reveal_type(A())  # type: A[int] # E: Revealed type is '__main__.A[builtins.int]'
+
+[case testUnionWithGenericTypeItemContext]
+from typing import TypeVar, Union, List
+
+T = TypeVar('T')
+
+def f(x: Union[T, List[int]]) -> Union[T, List[int]]: pass
+reveal_type(f(1)) # E: Revealed type is 'Union[builtins.int*, builtins.list[builtins.int]]'
+reveal_type(f([])) # E: Revealed type is 'builtins.list[builtins.int]'
+reveal_type(f(None)) # E: Revealed type is 'builtins.list[builtins.int]'
+[builtins fixtures/list.pyi]
+
+[case testUnionWithGenericTypeItemContextAndStrictOptional]
+# flags: --strict-optional
+from typing import TypeVar, Union, List
+
+T = TypeVar('T')
+
+def f(x: Union[T, List[int]]) -> Union[T, List[int]]: pass
+reveal_type(f(1)) # E: Revealed type is 'Union[builtins.int*, builtins.list[builtins.int]]'
+reveal_type(f([])) # E: Revealed type is 'builtins.list[builtins.int]'
+reveal_type(f(None)) # E: Revealed type is 'Union[builtins.None, builtins.list[builtins.int]]'
+[builtins fixtures/list.pyi]
+
+[case testUnionWithGenericTypeItemContextInMethod]
+from typing import TypeVar, Union, List, Generic
+
+T = TypeVar('T')
+S = TypeVar('S')
+
+class C(Generic[T]):
+    def f(self, x: Union[T, S]) -> Union[T, S]: pass
+
+c = C[List[int]]()
+reveal_type(c.f('')) # E: Revealed type is 'Union[builtins.list[builtins.int], builtins.str*]'
+reveal_type(c.f([1])) # E: Revealed type is 'builtins.list[builtins.int]'
+reveal_type(c.f([])) # E: Revealed type is 'builtins.list[builtins.int]'
+reveal_type(c.f(None)) # E: Revealed type is 'builtins.list[builtins.int]'
+[builtins fixtures/list.pyi]