PY-12002 Support union types for references in relationship fields

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()))

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);
+  }
 }