test fix: type is used for argument completion
authorIlya.Kazakevich <Ilya.Kazakevich@jetbrains.com>
Mon, 2 Nov 2015 19:28:31 +0000 (22:28 +0300)
committerIlya.Kazakevich <Ilya.Kazakevich@jetbrains.com>
Mon, 9 Nov 2015 17:33:00 +0000 (20:33 +0300)
python/src/com/jetbrains/python/psi/impl/references/KeywordArgumentCompletionUtil.java

index b31817cdf636954f0be38c7b81fad117a91bb3ee..1723d0c523a3b220739a9d12cf67811251a5ed02 100644 (file)
@@ -28,6 +28,7 @@ import com.jetbrains.python.psi.resolve.QualifiedResolveResult;
 import com.jetbrains.python.psi.search.PySuperMethodsSearch;
 import com.jetbrains.python.psi.types.*;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -40,9 +41,13 @@ public class KeywordArgumentCompletionUtil {
     if (callExpr != null) {
       PyExpression callee = callExpr.getCallee();
       if (callee instanceof PyReferenceExpression && element.getParent() == callExpr.getArgumentList()) {
-        final PyResolveContext resolveContext = PyResolveContext.defaultContext().withTypeEvalContext(context);
-        final QualifiedResolveResult result = ((PyReferenceExpression)callee).followAssignmentsChain(resolveContext);
-        PsiElement def = result.getElement();
+
+        PsiElement def = getElementByType(context, callee);
+        if (def == null) {
+          def = getElementByChain(context, (PyReferenceExpression)callee);
+        }
+
+
         if (def instanceof PyCallable) {
           addKeywordArgumentVariants((PyCallable)def, callExpr, ret);
         }
@@ -61,6 +66,24 @@ public class KeywordArgumentCompletionUtil {
     }
   }
 
+  @Nullable
+  private static PyElement getElementByType(@NotNull final TypeEvalContext context, @NotNull final PyExpression callee) {
+    final PyType pyType = context.getType(callee);
+    if (pyType instanceof PyFunctionType) {
+      return ((PyFunctionType)pyType).getCallable();
+    }
+    if (pyType instanceof PyClassType) {
+      return ((PyClassType)pyType).getPyClass();
+    }
+    return null;
+  }
+
+  private static PsiElement getElementByChain(@NotNull TypeEvalContext context, PyReferenceExpression callee) {
+    final PyResolveContext resolveContext = PyResolveContext.defaultContext().withTypeEvalContext(context);
+    final QualifiedResolveResult result = callee.followAssignmentsChain(resolveContext);
+    return result.getElement();
+  }
+
   private static void fetchCallablesFromUnion(@NotNull final List<LookupElement> ret,
                                               @NotNull final PyCallExpression callExpr,
                                               @NotNull final PyUnionType unionType,