From 326d440c1302c73e9fe4513bf2d7aee0167cc853 Mon Sep 17 00:00:00 2001
From: "Ilya.Kazakevich" <Ilya.Kazakevich@jetbrains.com>
Date: Mon, 10 Aug 2015 23:06:40 +0300
Subject: [PATCH] PY-3999 : Context passed to Keyword Argument Completion to
 allow better completion

---
 .../com/jetbrains/python/psi/impl/PyNamedParameterImpl.java  | 2 +-
 .../psi/impl/references/KeywordArgumentCompletionUtil.java   | 5 +++--
 .../python/psi/impl/references/PyReferenceImpl.java          | 3 ++-
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java b/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
index 5713741d97d8..cd604cf21500 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
@@ -219,7 +219,7 @@ public class PyNamedParameterImpl extends PyBaseElementImpl<PyNamedParameterStub
           final PyClass containingClass = func.getContainingClass();
           if (containingClass != null) {
             PyType initType = null;
-            final PyFunction init = containingClass.findInitOrNew(true, null);
+            final PyFunction init = containingClass.findInitOrNew(true, context);
             if (init != null && init != func) {
               initType = context.getReturnType(init);
               if (init.getContainingClass() != containingClass) {
diff --git a/python/src/com/jetbrains/python/psi/impl/references/KeywordArgumentCompletionUtil.java b/python/src/com/jetbrains/python/psi/impl/references/KeywordArgumentCompletionUtil.java
index 772c69e796f2..63df3d5e3c01 100644
--- a/python/src/com/jetbrains/python/psi/impl/references/KeywordArgumentCompletionUtil.java
+++ b/python/src/com/jetbrains/python/psi/impl/references/KeywordArgumentCompletionUtil.java
@@ -34,12 +34,13 @@ import java.util.HashSet;
 import java.util.List;
 
 public class KeywordArgumentCompletionUtil {
-  public static void collectFunctionArgNames(PyElement element, List<LookupElement> ret) {
+  public static void collectFunctionArgNames(PyElement element, List<LookupElement> ret,  final @NotNull TypeEvalContext context) {
     PyCallExpression callExpr = PsiTreeUtil.getParentOfType(element, PyCallExpression.class);
     if (callExpr != null) {
       PyExpression callee = callExpr.getCallee();
       if (callee instanceof PyReferenceExpression && element.getParent() == callExpr.getArgumentList()) {
-        final QualifiedResolveResult result = ((PyReferenceExpression)callee).followAssignmentsChain(PyResolveContext.defaultContext());
+        final PyResolveContext resolveContext = PyResolveContext.defaultContext().withTypeEvalContext(context);
+        final QualifiedResolveResult result = ((PyReferenceExpression)callee).followAssignmentsChain(resolveContext);
         PsiElement def = result.getElement();
         if (def instanceof PyFunction) {
           addKeywordArgumentVariants((PyFunction)def, callExpr, ret);
diff --git a/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java b/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java
index 68fea31df070..4f86c439cacf 100644
--- a/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java
@@ -578,8 +578,9 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
       PyResolveUtil.scopeCrawlUp(processor, owner, null, null);
     }
 
+    // This method is probably called for completion, so use appropriate context here
     // in a call, include function's arg names
-    KeywordArgumentCompletionUtil.collectFunctionArgNames(element, ret);
+    KeywordArgumentCompletionUtil.collectFunctionArgNames(element, ret, TypeEvalContext.codeCompletion(element.getProject(), element.getContainingFile()));
 
     // include builtin names
     final PyFile builtinsFile = PyBuiltinCache.getInstance(element).getBuiltinsFile();