From 65f44ea52256135274bbc3c71f7b773180cd61d8 Mon Sep 17 00:00:00 2001
From: "Ilya.Kazakevich" <Ilya.Kazakevich@jetbrains.com>
Date: Mon, 5 Oct 2015 21:19:56 +0300
Subject: [PATCH] PY-17052 Keyword arguments provided for unions with callables

---
 .../KeywordArgumentCompletionUtil.java        | 24 +++++++++++++++++--
 python/testData/callableArguments/file.after  |  3 +++
 python/testData/callableArguments/file.py     |  3 +++
 3 files changed, 28 insertions(+), 2 deletions(-)

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 0019f3f6203b..d4e7f9b1a7ba 100644
--- a/python/src/com/jetbrains/python/psi/impl/references/KeywordArgumentCompletionUtil.java
+++ b/python/src/com/jetbrains/python/psi/impl/references/KeywordArgumentCompletionUtil.java
@@ -26,7 +26,7 @@ import com.jetbrains.python.psi.impl.PyKeywordArgumentProvider;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.resolve.QualifiedResolveResult;
 import com.jetbrains.python.psi.search.PySuperMethodsSearch;
-import com.jetbrains.python.psi.types.TypeEvalContext;
+import com.jetbrains.python.psi.types.*;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.Collection;
@@ -51,6 +51,27 @@ public class KeywordArgumentCompletionUtil {
             addKeywordArgumentVariants(init, callExpr, ret);
           }
         }
+
+        final PyUnionType unionType = PyUtil.as(context.getType(callee), PyUnionType.class);
+        if (unionType != null) {
+          fetchCallablesFromUnion(ret, callExpr, unionType);
+        }
+      }
+    }
+  }
+
+  private static void fetchCallablesFromUnion(@NotNull final List<LookupElement> ret,
+                                              @NotNull final PyCallExpression callExpr,
+                                              @NotNull final PyUnionType unionType) {
+    for (final PyType memberType : unionType.getMembers()) {
+      if (memberType instanceof PyUnionType) {
+        fetchCallablesFromUnion(ret, callExpr, (PyUnionType)memberType);
+      }
+      if (memberType instanceof PyFunctionType) {
+        final PyFunctionType type = (PyFunctionType)memberType;
+        if (type.isCallable()) {
+          addKeywordArgumentVariants(type.getCallable(), callExpr, ret);
+        }
       }
     }
   }
@@ -75,7 +96,6 @@ public class KeywordArgumentCompletionUtil {
     else {
       addKeywordArgumentVariantsForCallable(callExpr, ret, parameters);
     }
-
   }
 
   private static void addKeywordArgumentVariantsForCallable(@NotNull final PyCallExpression callExpr,
diff --git a/python/testData/callableArguments/file.after b/python/testData/callableArguments/file.after
index a78f0bc8cdf7..6c3dd3f0656f 100644
--- a/python/testData/callableArguments/file.after
+++ b/python/testData/callableArguments/file.after
@@ -1,3 +1,6 @@
 spam = lambda eggs: eggs
 
+spam_union = 1 if 42 != 42 else spam
+
 spam(eggs=)
+spam_union(eggs=)
diff --git a/python/testData/callableArguments/file.py b/python/testData/callableArguments/file.py
index 257c2c714c13..744a49c4f9ae 100644
--- a/python/testData/callableArguments/file.py
+++ b/python/testData/callableArguments/file.py
@@ -1,3 +1,6 @@
 spam = lambda eggs: eggs
 
+spam_union = 1 if 42 != 42 else spam
+
 spam(egg<caret>)
+spam_union(egg<caret>)