1 files changed, 18 insertions, 3 deletions
diff --git a/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java
index 56249f1fdfa4..0dd75915de12 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java
@@ -27,6 +27,8 @@ import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Map;
+
 /**
  * @author yole
  */
@@ -61,10 +63,23 @@ public class PyLambdaExpressionImpl extends PyElementImpl implements PyLambdaExp
 
   @Nullable
   @Override
-  public PyType getReturnType(@NotNull TypeEvalContext context, @Nullable PyQualifiedExpression callSite) {
+  public PyType getReturnType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
     final PyExpression body = getBody();
-    if (body != null) return context.getType(body);
-    else return null;
+    return body != null ? context.getType(body) : null;
+  }
+
+  @Nullable
+  @Override
+  public PyType getCallType(@NotNull TypeEvalContext context, @NotNull PyQualifiedExpression callSite) {
+    return context.getReturnType(this);
+  }
+
+  @Nullable
+  @Override
+  public PyType getCallType(@Nullable PyExpression receiver,
+                            @NotNull Map<PyExpression, PyNamedParameter> parameters,
+                            @NotNull TypeEvalContext context) {
+    return context.getReturnType(this);
   }
 
   @Nullable