diff --git a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java index b787de24b388..7d196d6a5386 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java +++ b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java @@ -205,7 +205,7 @@ public class PyCallExpressionHelper { } } final List resolvedQualifiers = resolveResult != null ? resolveResult.getQualifiers() : null; - final List qualifiers = resolvedQualifiers != null ? resolvedQualifiers : Collections.emptyList(); + final List qualifiers = resolvedQualifiers != null ? resolvedQualifiers : Collections.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 callResultTypes = new HashSet(); + + 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 getCallTargetReturnType(@NotNull PyCallExpression call, @NotNull PsiElement target, @NotNull TypeEvalContext context) { @@ -686,7 +724,8 @@ public class PyCallExpressionHelper { final PositionalArgumentsAnalysisResults positionalResults = filterPositionalAndVariadicArguments(arguments); final List keywordArguments = filterKeywordArguments(arguments); final List variadicPositionalArguments = positionalResults.variadicPositionalArguments; - final Set positionalComponentsOfVariadicArguments = new LinkedHashSet(positionalResults.componentsOfVariadicPositionalArguments); + final Set positionalComponentsOfVariadicArguments = + new LinkedHashSet(positionalResults.componentsOfVariadicPositionalArguments); final List variadicKeywordArguments = filterVariadicKeywordArguments(arguments); final List allPositionalArguments = positionalResults.allPositionalArguments; diff --git a/python/testSrc/com/jetbrains/python/PyTypeTest.java b/python/testSrc/com/jetbrains/python/PyTypeTest.java index 32a318c62c9d..62d3578962a4 100644 --- a/python/testSrc/com/jetbrains/python/PyTypeTest.java +++ b/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" +