From 5dccf937e429c03be07dfb6ebf8355d1c9e6f8a7 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Mon, 20 Nov 2017 21:06:55 +0300
Subject: [PATCH] PY-12002 Support union types for references in relationship
 fields

---
 .../references/KeywordArgumentCompletionUtil.java  |  9 +--------
 .../com/jetbrains/python/psi/types/PyTypeUtil.java | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 8 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 a32e79012c89..9c8608ef97e8 100644
--- a/python/src/com/jetbrains/python/psi/impl/references/KeywordArgumentCompletionUtil.java
+++ b/python/src/com/jetbrains/python/psi/impl/references/KeywordArgumentCompletionUtil.java
@@ -51,14 +51,7 @@ public class KeywordArgumentCompletionUtil {
             calleeType = context.getType(implicit);
           }
         }
-        final StreamEx<PyType> types;
-        if (calleeType instanceof PyUnionType) {
-          types = StreamEx.of(((PyUnionType)calleeType).getMembers());
-        }
-        else {
-          types = StreamEx.of(calleeType);
-        }
-        final List<LookupElement> extra = types
+        final List<LookupElement> extra = PyTypeUtil.toStream(calleeType)
           .select(PyCallableType.class)
           .flatMap(type -> collectParameterNamesFromType(type, callExpr, context).stream())
           .map(name -> PyUtil.createNamedParameterLookup(name, element.getProject()))
diff --git a/python/src/com/jetbrains/python/psi/types/PyTypeUtil.java b/python/src/com/jetbrains/python/psi/types/PyTypeUtil.java
index 1f4714a79959..7ed8bf800839 100644
--- a/python/src/com/jetbrains/python/psi/types/PyTypeUtil.java
+++ b/python/src/com/jetbrains/python/psi/types/PyTypeUtil.java
@@ -19,6 +19,7 @@ import com.intellij.openapi.util.Key;
 import com.intellij.openapi.util.UserDataHolder;
 import com.intellij.psi.PsiElement;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
+import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -99,4 +100,17 @@ public final class PyTypeUtil {
       .map(dictClass -> new PyCollectionTypeImpl(dictClass, false, Arrays.asList(builtinCache.getStrType(), valueType)))
       .orElse(null);
   }
+
+  /**
+   * Given a type creates a stream of all its members if it's a union type or of only the type itself otherwise.
+   * <p>
+   * It allows to process types received as the result of multiresolve uniformly with the others.
+   */
+  @NotNull
+  public static StreamEx<PyType> toStream(@Nullable PyType type) {
+    if (type instanceof PyUnionType) {
+      return StreamEx.of(((PyUnionType)type).getMembers());
+    }
+    return StreamEx.of(type);
+  }
 }