From d438a0737b835be7a15cb9be583015c9e7c3edcc Mon Sep 17 00:00:00 2001 From: Semyon Proshev Date: Mon, 1 Aug 2016 18:40:38 +0300 Subject: [PATCH] Make PyTypeProvider.getCallType returns Ref instead of PyType. Attempt to make related methods more readable --- .../python/psi/impl/PyTypeProvider.java | 4 +- .../python/psi/types/PyTypeProviderBase.java | 94 +++++++++---------- .../NumpyDocStringTypeProvider.java | 45 +++++---- .../codeInsight/PyTypingTypeProvider.java | 20 ++-- .../stdlib/PyStdlibTypeProvider.java | 85 +++++++++-------- .../python/psi/impl/PyFunctionImpl.java | 25 +++-- .../jetbrains/python/pyi/PyiTypeProvider.java | 29 +++--- 7 files changed, 162 insertions(+), 140 deletions(-) diff --git a/python/psi-api/src/com/jetbrains/python/psi/impl/PyTypeProvider.java b/python/psi-api/src/com/jetbrains/python/psi/impl/PyTypeProvider.java index 3a74030dcb17..337fb6e23555 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/impl/PyTypeProvider.java +++ b/python/psi-api/src/com/jetbrains/python/psi/impl/PyTypeProvider.java @@ -1,5 +1,5 @@ /* - * Copyright 2000-2014 JetBrains s.r.o. + * Copyright 2000-2016 JetBrains s.r.o. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -43,7 +43,7 @@ public interface PyTypeProvider { Ref getReturnType(@NotNull PyCallable callable, @NotNull TypeEvalContext context); @Nullable - PyType getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context); + Ref getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context); @Nullable PyType getContextManagerVariableType(PyClass contextManager, PyExpression withExpression, TypeEvalContext context); diff --git a/python/psi-api/src/com/jetbrains/python/psi/types/PyTypeProviderBase.java b/python/psi-api/src/com/jetbrains/python/psi/types/PyTypeProviderBase.java index 07789a71c13d..ce4bc351c832 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/types/PyTypeProviderBase.java +++ b/python/psi-api/src/com/jetbrains/python/psi/types/PyTypeProviderBase.java @@ -1,5 +1,5 @@ /* - * Copyright 2000-2014 JetBrains s.r.o. + * Copyright 2000-2016 JetBrains s.r.o. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -17,6 +17,7 @@ package com.jetbrains.python.psi.types; import com.intellij.openapi.util.Ref; import com.intellij.psi.PsiElement; +import com.intellij.util.ObjectUtils; import com.intellij.util.containers.FactoryMap; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyTypeProvider; @@ -26,52 +27,18 @@ import org.jetbrains.annotations.Nullable; import java.util.Collection; import java.util.HashMap; import java.util.Map; +import java.util.Optional; /** * @author yole */ public class PyTypeProviderBase implements PyTypeProvider { - public PyTypeProviderBase() { - } - protected interface ReturnTypeCallback { - @Nullable - PyType getType(@Nullable PyCallSiteExpression callSite, @Nullable PyType qualifierType, TypeEvalContext context); - } - - private static class ReturnTypeDescriptor { - private final Map myStringToReturnTypeMap = new HashMap<>(); - - void put(String className, ReturnTypeCallback callback) { - myStringToReturnTypeMap.put(className, callback); - } - - @Nullable - public PyType get(PyFunction function, @Nullable PyCallSiteExpression callSite, TypeEvalContext context) { - PyClass containingClass = function.getContainingClass(); - if (containingClass != null) { - final ReturnTypeCallback typeCallback = myStringToReturnTypeMap.get(containingClass.getQualifiedName()); - if (typeCallback != null) { - final PyExpression callee = callSite instanceof PyCallExpression ? ((PyCallExpression)callSite).getCallee() : null; - final PyExpression qualifier = callee instanceof PyQualifiedExpression ? ((PyQualifiedExpression)callee).getQualifier() : null; - PyType qualifierType = qualifier != null ? context.getType(qualifier) : null; - return typeCallback.getType(callSite, qualifierType, context); - } - } - return null; - } - } - - private final ReturnTypeCallback mySelfTypeCallback = new ReturnTypeCallback() { - @Override - public PyType getType(@Nullable PyCallSiteExpression callSite, @Nullable PyType qualifierType, TypeEvalContext context) { - if (qualifierType instanceof PyClassType) { - PyClass aClass = ((PyClassType)qualifierType).getPyClass(); - return PyPsiFacade.getInstance(aClass.getProject()).createClassType(aClass, false); - } - return null; - } - }; + private final ReturnTypeCallback mySelfTypeCallback = (callSite, qualifierType, context) -> Optional + .ofNullable(ObjectUtils.tryCast(qualifierType, PyClassType.class)) + .map(PyClassType::getPyClass) + .map(pyClass -> PyPsiFacade.getInstance(pyClass.getProject()).createClassType(pyClass, false)) + .orElse(null); @SuppressWarnings({"MismatchedQueryAndUpdateOfCollection"}) private final Map myMethodToReturnTypeMap = new FactoryMap() { @@ -105,8 +72,8 @@ public class PyTypeProviderBase implements PyTypeProvider { @Nullable @Override - public PyType getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { - ReturnTypeDescriptor descriptor; + public Ref getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { + final ReturnTypeDescriptor descriptor; synchronized (myMethodToReturnTypeMap) { descriptor = myMethodToReturnTypeMap.get(function.getName()); } @@ -128,17 +95,50 @@ public class PyTypeProviderBase implements PyTypeProvider { return null; } - protected void registerSelfReturnType(String classQualifiedName, Collection methods) { + protected void registerSelfReturnType(@NotNull String classQualifiedName, @NotNull Collection methods) { registerReturnType(classQualifiedName, methods, mySelfTypeCallback); } - protected void registerReturnType(String classQualifiedName, - Collection methods, - final ReturnTypeCallback callback) { + protected void registerReturnType(@NotNull String classQualifiedName, + @NotNull Collection methods, + @NotNull ReturnTypeCallback callback) { synchronized (myMethodToReturnTypeMap) { for (String method : methods) { myMethodToReturnTypeMap.get(method).put(classQualifiedName, callback); } } } + + protected interface ReturnTypeCallback { + + @Nullable + PyType getType(@Nullable PyCallSiteExpression callSite, @Nullable PyType qualifierType, @NotNull TypeEvalContext context); + } + + private static class ReturnTypeDescriptor { + + private final Map myStringToReturnTypeMap = new HashMap<>(); + + public void put(@NotNull String classQualifiedName, @NotNull ReturnTypeCallback callback) { + myStringToReturnTypeMap.put(classQualifiedName, callback); + } + + @Nullable + public Ref get(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { + return Optional + .ofNullable(function.getContainingClass()) + .map(pyClass -> myStringToReturnTypeMap.get(pyClass.getQualifiedName())) + .map(typeCallback -> typeCallback.getType(callSite, getQualifierType(callSite, context), context)) + .map(Ref::create) + .orElse(null); + } + + @Nullable + private static PyType getQualifierType(@Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { + final PyExpression callee = callSite instanceof PyCallExpression ? ((PyCallExpression)callSite).getCallee() : null; + final PyExpression qualifier = callee instanceof PyQualifiedExpression ? ((PyQualifiedExpression)callee).getQualifier() : null; + + return qualifier != null ? context.getType(qualifier) : null; + } + } } diff --git a/python/src/com/jetbrains/numpy/codeInsight/NumpyDocStringTypeProvider.java b/python/src/com/jetbrains/numpy/codeInsight/NumpyDocStringTypeProvider.java index 7af729d7d377..25190e48cdb2 100644 --- a/python/src/com/jetbrains/numpy/codeInsight/NumpyDocStringTypeProvider.java +++ b/python/src/com/jetbrains/numpy/codeInsight/NumpyDocStringTypeProvider.java @@ -1,5 +1,5 @@ /* - * Copyright 2000-2014 JetBrains s.r.o. + * Copyright 2000-2016 JetBrains s.r.o. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -205,34 +205,34 @@ public class NumpyDocStringTypeProvider extends PyTypeProviderBase { @Nullable @Override - public PyType getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { + public Ref getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { if (isApplicable(function)) { final PyExpression callee = callSite instanceof PyCallExpression ? ((PyCallExpression)callSite).getCallee() : null; final NumpyDocString docString = forFunction(function, callee); if (docString != null) { final List returns = docString.getReturnFields(); final PyPsiFacade facade = getPsiFacade(function); + switch (returns.size()) { case 0: return null; case 1: // Function returns single value - final String typeName = returns.get(0).getType(); - if (StringUtil.isNotEmpty(typeName)) { - final PyType genericType = getPsiFacade(function).parseTypeAnnotation("T", function); - if (isUfuncType(function, typeName)) return genericType; - return parseNumpyDocType(function, typeName); - } - return null; + return Optional + .ofNullable(returns.get(0).getType()) + .filter(StringUtil::isNotEmpty) + .map(typeName -> isUfuncType(function, typeName) + ? facade.parseTypeAnnotation("T", function) + : parseNumpyDocType(function, typeName)) + .map(Ref::create) + .orElse(null); default: // Function returns a tuple - final ArrayList unionMembers = new ArrayList<>(); - + final List unionMembers = new ArrayList<>(); final List members = new ArrayList<>(); for (int i = 0; i < returns.size(); i++) { - SectionField ret = returns.get(i); - final String memberTypeName = ret.getType(); + final String memberTypeName = returns.get(i).getType(); final PyType returnType = StringUtil.isNotEmpty(memberTypeName) ? parseNumpyDocType(function, memberTypeName) : null; final boolean isOptional = StringUtil.isNotEmpty(memberTypeName) && memberTypeName.contains("optional"); @@ -250,13 +250,13 @@ public class NumpyDocStringTypeProvider extends PyTypeProviderBase { unionMembers.add(facade.createTupleType(members, function)); } } - if (unionMembers.isEmpty()) { - return facade.createTupleType(members, function); - } - return facade.createUnionType(unionMembers); + + final PyType type = unionMembers.isEmpty() ? facade.createTupleType(members, function) : facade.createUnionType(unionMembers); + return Ref.create(type); } } } + return null; } @@ -414,12 +414,9 @@ public class NumpyDocStringTypeProvider extends PyTypeProviderBase { @Nullable @Override public Ref getReturnType(@NotNull PyCallable callable, @NotNull TypeEvalContext context) { - if (callable instanceof PyFunction) { - final PyType type = getCallType((PyFunction)callable, null, context); - if (type != null) { - return Ref.create(type); - } - } - return null; + return Optional + .ofNullable(PyUtil.as(callable, PyFunction.class)) + .map(function -> getCallType(function, null, context)) + .orElse(null); } } diff --git a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java index 528203a7146b..8455e659e461 100644 --- a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java @@ -1,5 +1,5 @@ /* - * Copyright 2000-2014 JetBrains s.r.o. + * Copyright 2000-2016 JetBrains s.r.o. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -158,15 +158,17 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { @Nullable @Override - public PyType getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { - if ("typing.cast".equals(function.getQualifiedName()) && callSite instanceof PyCallExpression) { - final PyCallExpression callExpr = (PyCallExpression)callSite; - final PyExpression[] args = callExpr.getArguments(); - if (args.length > 0) { - final PyExpression typeExpr = args[0]; - return getType(typeExpr, new Context(context)); - } + public Ref getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { + if ("typing.cast".equals(function.getQualifiedName())) { + return Optional + .ofNullable(as(callSite, PyCallExpression.class)) + .map(PyCallExpression::getArguments) + .filter(args -> args.length > 0) + .map(args -> getType(args[0], new Context(context))) + .map(Ref::create) + .orElse(null); } + return null; } diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java index 7c64f2517723..04c72473709a 100644 --- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java @@ -17,6 +17,7 @@ package com.jetbrains.python.codeInsight.stdlib; import com.google.common.collect.ImmutableSet; import com.intellij.openapi.extensions.Extensions; +import com.intellij.openapi.util.Ref; import com.intellij.openapi.vfs.VirtualFile; import com.intellij.psi.PsiElement; import com.intellij.psi.util.QualifiedName; @@ -144,7 +145,16 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { @Nullable @Override - public PyType getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { + public Ref getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { + if (callSite != null && isListGetItem(function)) { + final PyExpression receiver = PyTypeChecker.getReceiver(callSite, function); + final Map mapping = PyCallExpressionHelper.mapArguments(callSite, function, context); + final Map substitutions = PyTypeChecker.unifyGenericCall(receiver, mapping, context); + if (substitutions != null) { + return analyzeListGetItemCallType(receiver, mapping, substitutions, context); + } + } + final String qname = getQualifiedName(function, callSite); if (qname != null) { if (OPEN_FUNCTIONS.contains(qname) && callSite instanceof PyCallExpression) { @@ -152,10 +162,7 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context); final PyCallExpression.PyArgumentsMapping mapping = callExpr.mapArguments(resolveContext); if (mapping.getMarkedCallee() != null) { - final PyType type = getOpenFunctionType(qname, mapping.getMappedParameters(), callSite); - if (type != null) { - return type; - } + return getOpenFunctionType(qname, mapping.getMappedParameters(), callSite); } } else if ("__builtin__.tuple.__add__".equals(qname) && callSite instanceof PyBinaryExpression) { @@ -164,15 +171,8 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { else if ("__builtin__.tuple.__mul__".equals(qname) && callSite instanceof PyBinaryExpression) { return getTupleMultiplicationResultType((PyBinaryExpression)callSite, context); } - else if (callSite != null && isListGetItem(function)) { - final PyExpression receiver = PyTypeChecker.getReceiver(callSite, function); - final Map mapping = PyCallExpressionHelper.mapArguments(callSite, function, context); - final Map substitutions = PyTypeChecker.unifyGenericCall(receiver, mapping, context); - if (substitutions != null) { - return analyzeListGetItemCallType(receiver, mapping, substitutions, context); - } - } } + return null; } @@ -186,10 +186,10 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { } @Nullable - private static PyType analyzeListGetItemCallType(@Nullable PyExpression receiver, - @NotNull Map parameters, - @NotNull Map substitutions, - @NotNull TypeEvalContext context) { + private static Ref analyzeListGetItemCallType(@Nullable PyExpression receiver, + @NotNull Map parameters, + @NotNull Map substitutions, + @NotNull TypeEvalContext context) { if (parameters.size() != 1 || substitutions.size() != 1) { return null; } @@ -204,25 +204,28 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { } if (PyABCUtil.isSubtype(firstArgumentType, PyNames.ABC_INTEGRAL, context)) { - return substitutions.values().iterator().next(); + return Ref.create(substitutions.values().iterator().next()); } if (PyNames.SLICE.equals(firstArgumentType.getName()) && firstArgumentType.isBuiltin()) { - return Optional - .ofNullable(receiver) - .map(context::getType) - .orElseGet(() -> PyTypeChecker.substitute(PyBuiltinCache.getInstance(receiver).getListType(), substitutions, context)); + return Ref.create( + Optional + .ofNullable(receiver) + .map(context::getType) + .orElseGet(() -> PyTypeChecker.substitute(PyBuiltinCache.getInstance(receiver).getListType(), substitutions, context)) + ); } return null; } @Nullable - private static PyType getTupleMultiplicationResultType(@NotNull PyBinaryExpression multiplication, @NotNull TypeEvalContext context) { + private static Ref getTupleMultiplicationResultType(@NotNull PyBinaryExpression multiplication, @NotNull TypeEvalContext context) { final PyTupleType leftTupleType = as(context.getType(multiplication.getLeftExpression()), PyTupleType.class); if (leftTupleType == null) { return null; } + PyExpression rightExpression = multiplication.getRightExpression(); if (rightExpression instanceof PyReferenceExpression) { final PsiElement target = ((PyReferenceExpression)rightExpression).getReference().resolve(); @@ -230,10 +233,12 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { rightExpression = ((PyTargetExpression)target).findAssignedValue(); } } + if (rightExpression instanceof PyNumericLiteralExpression && ((PyNumericLiteralExpression)rightExpression).isIntegerLiteral()) { if (leftTupleType.isHomogeneous()) { - return leftTupleType; + return Ref.create(leftTupleType); } + final int multiplier = ((PyNumericLiteralExpression)rightExpression).getBigIntegerValue().intValue(); final int originalSize = leftTupleType.getElementCount(); // Heuristic @@ -244,17 +249,19 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { elementTypes[i * originalSize + j] = leftTupleType.getElementType(j); } } - return PyTupleType.create(multiplication, elementTypes); + return Ref.create(PyTupleType.create(multiplication, elementTypes)); } } + return null; } @Nullable - private static PyType getTupleConcatenationResultType(@NotNull PyBinaryExpression addition, @NotNull TypeEvalContext context) { - final PyTupleType leftTupleType = as(context.getType(addition.getLeftExpression()), PyTupleType.class); + private static Ref getTupleConcatenationResultType(@NotNull PyBinaryExpression addition, @NotNull TypeEvalContext context) { if (addition.getRightExpression() != null) { + final PyTupleType leftTupleType = as(context.getType(addition.getLeftExpression()), PyTupleType.class); final PyTupleType rightTupleType = as(context.getType(addition.getRightExpression()), PyTupleType.class); + if (leftTupleType != null && rightTupleType != null) { if (leftTupleType.isHomogeneous() || rightTupleType.isHomogeneous()) { // We may try to find the common type of elements of two homogeneous tuple as an alternative @@ -268,9 +275,11 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { for (int i = 0; i < rightTupleType.getElementCount(); i++) { elementTypes[i + leftTupleType.getElementCount()] = rightTupleType.getElementType(i); } - return PyTupleType.create(addition, elementTypes); + + return Ref.create(PyTupleType.create(addition, elementTypes)); } } + return null; } @@ -310,10 +319,10 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { return null; } - @Nullable - private static PyType getOpenFunctionType(@NotNull String callQName, - @NotNull Map arguments, - @NotNull PsiElement anchor) { + @NotNull + private static Ref getOpenFunctionType(@NotNull String callQName, + @NotNull Map arguments, + @NotNull PsiElement anchor) { String mode = "r"; for (Map.Entry entry : arguments.entrySet()) { final PyNamedParameter parameter = entry.getValue(); @@ -328,17 +337,17 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { } } } - final LanguageLevel level = LanguageLevel.forElement(anchor); - if (level.isPy3K() || "io.open".equals(callQName)) { + if (LanguageLevel.forElement(anchor).isAtLeast(LanguageLevel.PYTHON30) || "io.open".equals(callQName)) { if (mode.contains("b")) { - return PyTypeParser.getTypeByName(anchor, PY3K_BINARY_FILE_TYPE); - } else { - return PyTypeParser.getTypeByName(anchor, PY3K_TEXT_FILE_TYPE); + return Ref.create(PyTypeParser.getTypeByName(anchor, PY3K_BINARY_FILE_TYPE)); + } + else { + return Ref.create(PyTypeParser.getTypeByName(anchor, PY3K_TEXT_FILE_TYPE)); } } - return PyTypeParser.getTypeByName(anchor, PY2K_FILE_TYPE); + return Ref.create(PyTypeParser.getTypeByName(anchor, PY2K_FILE_TYPE)); } @Nullable diff --git a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java index dca8496cb54f..e88304ef287a 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java @@ -194,13 +194,10 @@ public class PyFunctionImpl extends PyBaseElementImpl implements for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) { final Ref returnTypeRef = typeProvider.getReturnType(this, context); if (returnTypeRef != null) { - final PyType returnType = returnTypeRef.get(); - if (returnType != null) { - returnType.assertValid(typeProvider.toString()); - } - return returnType; + return derefType(returnTypeRef, typeProvider); } } + if (context.allowReturnTypes(this)) { final Ref yieldTypeRef = getYieldStatementType(context); if (yieldTypeRef != null) { @@ -208,6 +205,7 @@ public class PyFunctionImpl extends PyBaseElementImpl implements } return getReturnStatementType(context); } + return null; } @@ -215,17 +213,26 @@ public class PyFunctionImpl extends PyBaseElementImpl implements @Override public PyType getCallType(@NotNull TypeEvalContext context, @NotNull PyCallSiteExpression callSite) { for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) { - final PyType type = typeProvider.getCallType(this, callSite, context); - if (type != null) { - type.assertValid(typeProvider.toString()); - return type; + final Ref typeRef = typeProvider.getCallType(this, callSite, context); + if (typeRef != null) { + return derefType(typeRef, typeProvider); } } + final PyExpression receiver = PyTypeChecker.getReceiver(callSite, this); final Map mapping = PyCallExpressionHelper.mapArguments(callSite, this, context); return getCallType(receiver, mapping, context); } + @Nullable + private static PyType derefType(@NotNull Ref typeRef, @NotNull PyTypeProvider typeProvider) { + final PyType type = typeRef.get(); + if (type != null) { + type.assertValid(typeProvider.toString()); + } + return type; + } + @Nullable @Override public PyType getCallType(@Nullable PyExpression receiver, diff --git a/python/src/com/jetbrains/python/pyi/PyiTypeProvider.java b/python/src/com/jetbrains/python/pyi/PyiTypeProvider.java index 923b1f6521db..dc7cf7eb5745 100644 --- a/python/src/com/jetbrains/python/pyi/PyiTypeProvider.java +++ b/python/src/com/jetbrains/python/pyi/PyiTypeProvider.java @@ -91,42 +91,49 @@ public class PyiTypeProvider extends PyTypeProviderBase { @Nullable @Override - public PyType getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { + public Ref getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { if (callSite != null) { final PsiElement pythonStub = PyiUtil.getPythonStub(function); + if (pythonStub instanceof PyFunction) { - final PyFunction functionStub = (PyFunction)pythonStub; - return getOverloadedCallType(functionStub, callSite, context); + return getOverloadedCallType((PyFunction)pythonStub, callSite, context); } else if (function.getContainingFile() instanceof PyiFile) { return getOverloadedCallType(function, callSite, context); } } + return null; } @Nullable - private static PyType getOverloadedCallType(@NotNull PyFunction function, @NotNull PyCallSiteExpression callSite, - @NotNull TypeEvalContext context) { + private static Ref getOverloadedCallType(@NotNull PyFunction function, + @NotNull PyCallSiteExpression callSite, + @NotNull TypeEvalContext context) { if (isOverload(function, context)) { - final List matchedReturnTypes = new ArrayList<>(); - final List allReturnTypes = new ArrayList<>(); final List overloads = getOverloads(function, context); + final List allReturnTypes = new ArrayList<>(); + final List matchedReturnTypes = new ArrayList<>(); + for (PyFunction overload : overloads) { - final Map mapping = mapArguments(callSite, overload, context); - final PyExpression receiver = PyTypeChecker.getReceiver(callSite, overload); - final Map substitutions = PyTypeChecker.unifyGenericCall(receiver, mapping, context); final PyType returnType = context.getReturnType(overload); if (!PyTypeChecker.hasGenerics(returnType, context)) { allReturnTypes.add(returnType); } + + final PyExpression receiver = PyTypeChecker.getReceiver(callSite, overload); + final Map mapping = mapArguments(callSite, overload, context); + final Map substitutions = PyTypeChecker.unifyGenericCall(receiver, mapping, context); + final PyType unifiedType = substitutions != null ? PyTypeChecker.substitute(returnType, substitutions, context) : null; if (unifiedType != null) { matchedReturnTypes.add(unifiedType); } } - return PyUnionType.union(matchedReturnTypes.isEmpty() ? allReturnTypes : matchedReturnTypes); + + return Ref.create(PyUnionType.union(matchedReturnTypes.isEmpty() ? allReturnTypes : matchedReturnTypes)); } + return null; }