PY-17148 union with callable types should also be callable

python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java

@@ -205,7 +205,7 @@ public class PyCallExpressionHelper {
       }
     }
     final List<PyExpression> resolvedQualifiers = resolveResult != null ? resolveResult.getQualifiers() : null;
-    final List<PyExpression> qualifiers = resolvedQualifiers != null ?  resolvedQualifiers : Collections.<PyExpression>emptyList();
+    final List<PyExpression> qualifiers = resolvedQualifiers != null ? resolvedQualifiers : Collections.<PyExpression>emptyList();
     final TypeEvalContext context = resolveContext.getTypeEvalContext();
     if (resolved instanceof PyFunction) {
       final PyFunction function = (PyFunction)resolved;
@@ -254,8 +254,8 @@ public class PyCallExpressionHelper {
    * Calls the {@link #getImplicitArgumentCount(PyCallable, PyFunction.Modifier, boolean, boolean, boolean)} full version}
    * with null flags and with isByInstance inferred directly from call site (won't work with reassigned bound methods).
    *
-   * @param callReference       the call site, where arguments are given.
-   * @param function resolved method which is being called; plain functions are OK but make little sense.
+   * @param callReference the call site, where arguments are given.
+   * @param function      resolved method which is being called; plain functions are OK but make little sense.
    * @return a non-negative number of parameters that are implicit to this call.
    */
   public static int getImplicitArgumentCount(@NotNull final PyReferenceExpression callReference, @NotNull PyFunction function,
@@ -481,6 +481,9 @@ public class PyCallExpressionHelper {
           final PyCallableType callableType = (PyCallableType)type;
           return callableType.getCallType(context, call);
         }
+        if (type instanceof PyUnionType) {
+          return getCallResultTypeFromUnion(call, context, (PyUnionType)type);
+        }
         return null;
       }
     }
@@ -489,6 +492,41 @@ public class PyCallExpressionHelper {
     }
   }
 
+  /**
+   * @return type that union will return if you call it
+   */
+  @Nullable
+  private static PyType getCallResultTypeFromUnion(@NotNull final PyCallSiteExpression call,
+                                                   @NotNull final TypeEvalContext context,
+                                                   @NotNull final PyUnionType type) {
+    final Collection<PyType> callResultTypes = new HashSet<PyType>();
+
+    for (final PyType memberType : type.getMembers()) {
+      final Boolean callable = PyTypeChecker.isCallable(memberType);
+      if (!((callable != null && callable && memberType instanceof PyCallableType))) {
+        continue;
+      }
+      final PyCallableType callableMemberType = (PyCallableType)memberType;
+
+      if (!callableMemberType.isCallable()) {
+        continue;
+      }
+      final PyType callResultType = callableMemberType.getCallType(context, call);
+      if (callResultType != null) {
+        callResultTypes.add(callResultType);
+      }
+    }
+
+
+    if (callResultTypes.size() == 1) {
+      return callResultTypes.iterator().next();
+    }
+    if (callResultTypes.size() > 1) {
+      return PyUnionType.union(callResultTypes);
+    }
+    return null;
+  }
+
   @Nullable
   private static Ref<? extends PyType> getCallTargetReturnType(@NotNull PyCallExpression call, @NotNull PsiElement target,
                                                                @NotNull TypeEvalContext context) {
@@ -686,7 +724,8 @@ public class PyCallExpressionHelper {
     final PositionalArgumentsAnalysisResults positionalResults = filterPositionalAndVariadicArguments(arguments);
     final List<PyKeywordArgument> keywordArguments = filterKeywordArguments(arguments);
     final List<PyExpression> variadicPositionalArguments = positionalResults.variadicPositionalArguments;
-    final Set<PyExpression> positionalComponentsOfVariadicArguments = new LinkedHashSet<PyExpression>(positionalResults.componentsOfVariadicPositionalArguments);
+    final Set<PyExpression> positionalComponentsOfVariadicArguments =
+      new LinkedHashSet<PyExpression>(positionalResults.componentsOfVariadicPositionalArguments);
     final List<PyExpression> variadicKeywordArguments = filterVariadicKeywordArguments(arguments);
 
     final List<PyExpression> allPositionalArguments = positionalResults.allPositionalArguments;

python/testSrc/com/jetbrains/python/PyTypeTest.java

@@ -29,6 +29,23 @@ import java.util.List;
  * @author yole
  */
 public class PyTypeTest extends PyTestCase {
+  /**
+   * Call of union returns union of all callable types in this union
+   */
+  public void testCallableInUnion() throws Exception {
+    doTest("str",
+           "import random\n" +
+           "def spam():\n" +
+           "    return \"D\"\n" +
+           "class Eggs:\n" +
+           "    pass\n" +
+           "class Eggs2:\n" +
+           "    pass\n" +
+           "dd = spam if random.randint != 42 else Eggs2()\n" +
+           "var = dd if random.randint != 42 else dd\n" +
+           "expr = var()");
+  }
+
   public void testTupleType() {
     doTest("str",
            "t = ('a', 2)\n" +