From dd898e2d02b714c3f5f933c5c60d0c8524c1bf7e Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Mon, 3 Mar 2014 14:56:50 +0400
Subject: [PATCH] Introduced Callable.getCallType() instead of
 PyFunctionImpl.getReturnTypeWithoutCallSite()

---
 .../com/jetbrains/python/psi/Callable.java    | 10 +++++
 .../python/psi/impl/PyFunctionImpl.java       | 44 ++++++++-----------
 .../psi/impl/PyLambdaExpressionImpl.java      | 10 +++++
 .../psi/impl/PyTargetExpressionImpl.java      | 10 ++---
 4 files changed, 43 insertions(+), 31 deletions(-)

diff --git a/python/psi-api/src/com/jetbrains/python/psi/Callable.java b/python/psi-api/src/com/jetbrains/python/psi/Callable.java
index b979b0deb79a..7eb9727ebbb4 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/Callable.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/Callable.java
@@ -20,6 +20,8 @@ import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Map;
+
 /**
  * Something that can be called, passed parameters to, and return something back.
 
@@ -45,6 +47,14 @@ public interface Callable extends PyTypedElement, PyQualifiedNameOwner {
   @Nullable
   PyType getCallType(@NotNull TypeEvalContext context, @Nullable PyQualifiedExpression callSite);
 
+  /**
+   * Returns the type of the call to the callable where the call site is specified by the optional receiver and the arguments to parameters
+   * mapping.
+   */
+  @Nullable
+  PyType getCallType(@Nullable PyExpression receiver, @NotNull Map<PyExpression, PyNamedParameter> parameters,
+                     @NotNull TypeEvalContext context);
+
   /**
    * @return a methods returns itself, non-method callables return null.
    */
diff --git a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
index 5a5c074c225e..62776f013af1 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
@@ -15,7 +15,6 @@
  */
 package com.jetbrains.python.psi.impl;
 
-import com.google.common.collect.Maps;
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.util.Pair;
@@ -235,36 +234,25 @@ public class PyFunctionImpl extends PyPresentableElementImpl<PyFunctionStub> imp
       return type;
     }
     final PyTypeChecker.AnalyzeCallResults results = PyTypeChecker.analyzeCallSite(callSite, context);
-    if (PyTypeChecker.hasGenerics(type, context)) {
-      if (results != null) {
-        final Map<PyGenericType, PyType> substitutions = PyTypeChecker.unifyGenericCall(results.getReceiver(), results.getArguments(),
-                                                                                        context);
-        type = substitutions != null ? PyTypeChecker.substitute(type, substitutions, context) : null;
-      }
-      else {
-        type = null;
-      }
-    }
     if (results != null) {
-      type = replaceSelf(type, results.getReceiver(), context);
-    }
-    if (results != null && isDynamicallyEvaluated(results.getArguments().values(), context)) {
-      return PyUnionType.createWeakType(type);
+      return analyzeCallType(type, results.getReceiver(), results.getArguments(), context);
     }
     return type;
   }
 
   @Nullable
-  /**
-   * Suits when there is no call site(e.g. implicit __iter__ call in statement for)
-   */
-  public PyType getReturnTypeWithoutCallSite(@NotNull TypeEvalContext context,
-                                             @Nullable PyExpression receiver) {
-    PyType type = getReturnType(context);
+  @Override
+  public PyType getCallType(@Nullable PyExpression receiver,
+                            @NotNull Map<PyExpression, PyNamedParameter> parameters,
+                            @NotNull TypeEvalContext context) {
+    return analyzeCallType(getReturnType(context), receiver, parameters, context);
+  }
+
+  @Nullable
+  private PyType analyzeCallType(@Nullable PyType type, @Nullable PyExpression receiver,
+                                 @NotNull Map<PyExpression, PyNamedParameter> parameters, @NotNull TypeEvalContext context) {
     if (PyTypeChecker.hasGenerics(type, context)) {
-      final Map<PyGenericType, PyType> substitutions = PyTypeChecker.unifyGenericCall(receiver,
-                                                                                      Maps.<PyExpression, PyNamedParameter>newHashMap(),
-                                                                                      context);
+      final Map<PyGenericType, PyType> substitutions = PyTypeChecker.unifyGenericCall(receiver, parameters, context);
       if (substitutions != null) {
         type = PyTypeChecker.substitute(type, substitutions, context);
       }
@@ -272,7 +260,13 @@ public class PyFunctionImpl extends PyPresentableElementImpl<PyFunctionStub> imp
         type = null;
       }
     }
-    return replaceSelf(type, receiver, context);
+    if (receiver != null) {
+      type = replaceSelf(type, receiver, context);
+    }
+    if (type != null && isDynamicallyEvaluated(parameters.values(), context)) {
+      type = PyUnionType.createWeakType(type);
+    }
+    return type;
   }
 
   @Nullable
diff --git a/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java
index f39a88947adb..97f1303cb24e 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java
@@ -28,6 +28,8 @@ import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Map;
+
 /**
  * @author yole
  */
@@ -82,6 +84,14 @@ public class PyLambdaExpressionImpl extends PyElementImpl implements PyLambdaExp
     return getReturnType(context);
   }
 
+  @Nullable
+  @Override
+  public PyType getCallType(@Nullable PyExpression receiver,
+                            @NotNull Map<PyExpression, PyNamedParameter> parameters,
+                            @NotNull TypeEvalContext context) {
+    return getReturnType(context);
+  }
+
   @Nullable
   public PyExpression getBody() {
     return PsiTreeUtil.getChildOfType(this, PyExpression.class);
diff --git a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
index 3ce0c51e7183..060c39ccdee0 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
@@ -54,6 +54,7 @@ import org.jetbrains.annotations.Nullable;
 import javax.swing.*;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 
 /**
@@ -243,8 +244,8 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
       if (exprType instanceof PyClassType) {
         final PyClass cls = ((PyClassType)exprType).getPyClass();
         final PyFunction enter = cls.findMethodByName(PyNames.ENTER, true);
-        if (enter instanceof PyFunctionImpl) {
-          final PyType enterType = ((PyFunctionImpl)enter).getReturnTypeWithoutCallSite(context, expression);
+        if (enter != null) {
+          final PyType enterType = enter.getCallType(expression, Collections.<PyExpression, PyNamedParameter>emptyMap(), context);
           if (enterType != null) {
             return enterType;
           }
@@ -438,10 +439,7 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
   @Nullable
   private static PyType getContextSensitiveType(@NotNull PyFunction function, @NotNull TypeEvalContext context,
                                                 @Nullable PyExpression source) {
-    if (function instanceof PyFunctionImpl) {
-      return ((PyFunctionImpl)function).getReturnTypeWithoutCallSite(context, source);
-    }
-    return function.getCallType(context, null);
+    return function.getCallType(source, Collections.<PyExpression, PyNamedParameter>emptyMap(), context);
   }
 
   @Nullable