Fix lambda inference with simple Union contexts (#2630)

Author: ddfisher

mypy/checkexpr.py

@@ -1772,6 +1772,12 @@ def infer_lambda_type_using_context(self, e: FuncExpr) -> CallableType:
         """
         # TODO also accept 'Any' context
         ctx = self.chk.type_context[-1]
+
+        if isinstance(ctx, UnionType):
+            callables = [t for t in ctx.items if isinstance(t, CallableType)]
+            if len(callables) == 1:
+                ctx = callables[0]
+
         if not ctx or not isinstance(ctx, CallableType):
             return None
 

test-data/unit/check-optional.test

@@ -518,3 +518,10 @@ if x is not None:
     x.x = 1 # OK here
 
 [builtins fixtures/ops.pyi]
+
+[case testOptionalLambdaInference]
+from typing import Optional, Callable
+f = None # type: Optional[Callable[[int], None]]
+f = lambda x: None
+f(0)
+[builtins fixtures/function.pyi]