diff --git a/python/python-psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java b/python/python-psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java index 043bd7ae343e..0910ad29eab3 100644 --- a/python/python-psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java +++ b/python/python-psi-api/src/com/jetbrains/python/psi/stubs/PyClassStub.java @@ -43,7 +43,7 @@ public interface PyClassStub extends NamedStub, PyVersionSpecificStub { * @return literal text of expressions in the base classes list. */ @NotNull - List getSuperClassesText(); + List<@NotNull String> getSuperClassesText(); @ApiStatus.Internal diff --git a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.kt b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.kt index 7511894a0a2a..2c6bfdfa6359 100644 --- a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.kt +++ b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.kt @@ -31,7 +31,7 @@ import org.jetbrains.annotations.ApiStatus import java.util.concurrent.ConcurrentMap import kotlin.concurrent.Volatile -open class TypeEvalContext private constructor( +sealed class TypeEvalContext( /** * @return context constraints (see [TypeEvalConstraints] */ @@ -267,7 +267,7 @@ open class TypeEvalContext private constructor( get() = constraints.myOrigin @ApiStatus.Internal - fun getContextTypeCache(): Map, PyType?> { + fun getContextTypeCache(): MutableMap, PyType?> { return contextTypeCache } @@ -289,7 +289,7 @@ open class TypeEvalContext private constructor( isSameVirtualFile(constraints.myOrigin, getContextFile(element)) } - object PyNullType : PyType { + private object PyNullType : PyType { override fun resolveMember( name: String, location: PyExpression?, @@ -315,6 +315,9 @@ open class TypeEvalContext private constructor( } } + private class TypeEvalContextImpl(allowDataFlow: Boolean, allowStubToAST: Boolean, allowCallContext: Boolean, origin: PsiFile?) : + TypeEvalContext(allowDataFlow, allowStubToAST, allowCallContext, origin) + private class AssumptionContext(val myParent: TypeEvalContext, element: PyTypedElement, type: PyType?) : TypeEvalContext(myParent.constraints) { init { @@ -418,6 +421,7 @@ open class TypeEvalContext private constructor( } } + @ApiStatus.Internal companion object { private fun getConcurrentMapForCaching(): ConcurrentMap { // In the current implementation, this value is only used to initialize the map and is basically ignored @@ -451,7 +455,7 @@ open class TypeEvalContext private constructor( */ @JvmStatic fun codeCompletion(project: Project, origin: PsiFile?): TypeEvalContext { - return getContextFromCache(project, TypeEvalContext(true, true, true, origin)) + return getContextFromCache(project, TypeEvalContextImpl(true, true, true, origin)) } /** @@ -465,7 +469,7 @@ open class TypeEvalContext private constructor( */ @JvmStatic fun userInitiated(project: Project, origin: PsiFile?): TypeEvalContext { - return getContextFromCache(project, TypeEvalContext(true, true, false, origin)) + return getContextFromCache(project, TypeEvalContextImpl(true, true, false, origin)) } /** @@ -489,7 +493,7 @@ open class TypeEvalContext private constructor( */ @JvmStatic fun codeInsightFallback(project: Project?): TypeEvalContext { - val anchor = TypeEvalContext(false, false, false, null) + val anchor = TypeEvalContextImpl(false, false, false, null) if (project != null) { return getContextFromCache(project, anchor) } @@ -504,14 +508,14 @@ open class TypeEvalContext private constructor( */ @JvmStatic fun deepCodeInsight(project: Project): TypeEvalContext { - return getContextFromCache(project, TypeEvalContext(false, true, false, null)) + return getContextFromCache(project, TypeEvalContextImpl(false, true, false, null)) } private fun buildCodeAnalysisContext(origin: PsiFile?): TypeEvalContext { if (Registry.`is`("python.optimized.type.eval.context")) { return OptimizedTypeEvalContext(false, false, false, origin) } - return TypeEvalContext(false, false, false, origin) + return TypeEvalContextImpl(false, false, false, origin) } /** diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyAncestorTypeProvider.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyAncestorTypeProvider.kt index 06cebe4ec50f..313a8ab9c64a 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyAncestorTypeProvider.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyAncestorTypeProvider.kt @@ -3,7 +3,6 @@ package com.jetbrains.python.codeInsight.typing import com.intellij.openapi.util.Ref import com.jetbrains.python.PyNames -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.getReturnTypeAnnotation import com.jetbrains.python.psi.PyCallable import com.jetbrains.python.psi.PyFunction import com.jetbrains.python.psi.PyNamedParameter @@ -87,7 +86,7 @@ private fun getReturnTypeFromSupertype(function: PyFunction, context: TypeEvalCo val overriddenFunction = getOverriddenFunction(function, context) if (overriddenFunction != null) { - val superFunctionAnnotation = getReturnTypeAnnotation(overriddenFunction, context) + val superFunctionAnnotation = PyTypingTypeProvider.getReturnTypeAnnotation(overriddenFunction, context) if (superFunctionAnnotation != null) { val typeRef = PyTypingTypeProvider.getType(superFunctionAnnotation, context) if (typeRef != null && function.isAsync == overriddenFunction.isAsync) { diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyProtocols.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyProtocols.kt index 4e85768eadd9..fa6f2625ba5b 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyProtocols.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyProtocols.kt @@ -3,8 +3,6 @@ package com.jetbrains.python.codeInsight.typing import com.intellij.psi.util.contextOfType import com.jetbrains.python.PyNames -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.PROTOCOL -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.PROTOCOL_EXT import com.jetbrains.python.psi.PyClass import com.jetbrains.python.psi.PyFunction import com.jetbrains.python.psi.PyKnownDecorator.TYPING_RUNTIME @@ -107,5 +105,5 @@ fun inspectProtocolSubclass(protocol: PyClassType, subclass: PyClassType, contex private fun containsProtocol(types: List) = types.any { type -> val classQName = type?.classQName - PROTOCOL == classQName || PROTOCOL_EXT == classQName + PyTypingTypeProvider.PROTOCOL == classQName || PyTypingTypeProvider.PROTOCOL_EXT == classQName } diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypeHintProvider.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypeHintProvider.kt index af983cbaf5d2..808e4da3019b 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypeHintProvider.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypeHintProvider.kt @@ -13,12 +13,12 @@ import org.jetbrains.annotations.ApiStatus.Experimental @Experimental @ApiStatus.Internal interface PyTypeHintProvider { - fun parseTypeHint(typeHint: PyExpression, alias: PyQualifiedNameOwner?, resolved: PsiElement, context: TypeEvalContext): Ref? + fun parseTypeHint(typeHint: PyExpression, alias: PyQualifiedNameOwner?, resolved: PsiElement, context: TypeEvalContext): Ref? companion object { private val EP_NAME: ExtensionPointName = ExtensionPointName.create("Pythonid.typeHintProvider"); - fun parseTypeHint(typeHint: PyExpression, alias: PyQualifiedNameOwner?, resolved: PsiElement, context: TypeEvalContext): Ref? { + fun parseTypeHint(typeHint: PyExpression, alias: PyQualifiedNameOwner?, resolved: PsiElement, context: TypeEvalContext): Ref? { return EP_NAME.extensionList.firstNotNullOfOrNull { it.parseTypeHint(typeHint, alias, resolved, context) } } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypedDictTypeProvider.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypedDictTypeProvider.kt index 64e5b2cb0468..6799956b35b5 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypedDictTypeProvider.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypedDictTypeProvider.kt @@ -6,17 +6,6 @@ import com.intellij.psi.PsiElement import com.intellij.psi.impl.source.resolve.FileContextUtil import com.jetbrains.python.PyCustomType import com.jetbrains.python.PyNames -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.MAPPING_GET -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.NOT_REQUIRED -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.NOT_REQUIRED_EXT -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.READONLY -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.READONLY_EXT -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.REQUIRED -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.REQUIRED_EXT -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.TYPED_DICT -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.TYPED_DICT_EXT -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.getType -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.resolveToQualifiedNames import com.jetbrains.python.psi.LanguageLevel import com.jetbrains.python.psi.PyBoolLiteralExpression import com.jetbrains.python.psi.PyCallExpression @@ -72,15 +61,16 @@ class PyTypedDictTypeProvider : PyTypeProviderBase() { } companion object { - val nameIsTypedDict = { name: String? -> name == TYPED_DICT || name == TYPED_DICT_EXT } + val nameIsTypedDict = { name: String? -> name == PyTypingTypeProvider.TYPED_DICT || name == PyTypingTypeProvider.TYPED_DICT_EXT } fun isGetMethodToOverride(call: PyCallExpression, context: TypeEvalContext): Boolean { val callee = call.callee - return callee != null && resolveToQualifiedNames(callee, context).any { it == "dict.get" /* py3 */ || it == MAPPING_GET /* py2 */ } + return callee != null && PyTypingTypeProvider.resolveToQualifiedNames(callee, context) + .any { it == "dict.get" /* py3 */ || it == PyTypingTypeProvider.MAPPING_GET /* py2 */ } } fun isTypedDict(expression: PyExpression, context: TypeEvalContext): Boolean { - return resolveToQualifiedNames(expression, context).any(nameIsTypedDict) + return PyTypingTypeProvider.resolveToQualifiedNames(expression, context).any(nameIsTypedDict) } fun isTypingTypedDictInheritor(cls: PyClass, context: TypeEvalContext): Boolean { @@ -266,14 +256,14 @@ class PyTypedDictTypeProvider : PyTypeProviderBase() { val result = mutableListOf() expression.accept(object : PyRecursiveElementVisitor() { override fun visitPySubscriptionExpression(node: PySubscriptionExpression) { - val resolvedNames = resolveToQualifiedNames(node.operand, context) - if (resolvedNames.any { name -> REQUIRED == name || REQUIRED_EXT == name }) { + val resolvedNames = PyTypingTypeProvider.resolveToQualifiedNames(node.operand, context) + if (resolvedNames.any { name -> PyTypingTypeProvider.REQUIRED == name || PyTypingTypeProvider.REQUIRED_EXT == name }) { result.add(TypedDictFieldQualifier.REQUIRED) } - else if (resolvedNames.any { name -> NOT_REQUIRED == name || NOT_REQUIRED_EXT == name }) { + else if (resolvedNames.any { name -> PyTypingTypeProvider.NOT_REQUIRED == name || PyTypingTypeProvider.NOT_REQUIRED_EXT == name }) { result.add(TypedDictFieldQualifier.NOT_REQUIRED) } - else if (resolvedNames.any { name -> READONLY == name || READONLY_EXT == name }) { + else if (resolvedNames.any { name -> PyTypingTypeProvider.READONLY == name || PyTypingTypeProvider.READONLY_EXT == name }) { result.add(TypedDictFieldQualifier.READ_ONLY) } super.visitPySubscriptionExpression(node) @@ -367,7 +357,7 @@ class PyTypedDictTypeProvider : PyTypeProviderBase() { if (expr is PySubscriptionExpression) { qualifiers = parseTypedDictFieldQualifiers(expr, context) } - return if (expr != null) Pair(getType(expr, context), qualifiers) else null + return if (expr != null) Pair(PyTypingTypeProvider.getType(expr, context), qualifiers) else null } } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt index dc55769668b0..efe75a83e6e5 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt @@ -1,602 +1,301 @@ // Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. -package com.jetbrains.python.codeInsight.typing; +package com.jetbrains.python.codeInsight.typing -import com.dynatrace.hash4j.hashing.HashValue128; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import com.intellij.openapi.application.PathManager; -import com.intellij.openapi.project.Project; -import com.intellij.openapi.util.Couple; -import com.intellij.openapi.util.Key; -import com.intellij.openapi.util.NlsSafe; -import com.intellij.openapi.util.Pair; -import com.intellij.openapi.util.Ref; -import com.intellij.openapi.util.TextRange; -import com.intellij.psi.PsiComment; -import com.intellij.psi.PsiElement; -import com.intellij.psi.PsiFile; -import com.intellij.psi.ResolveResult; -import com.intellij.psi.impl.source.resolve.FileContextUtil; -import com.intellij.psi.search.GlobalSearchScope; -import com.intellij.psi.util.CachedValueProvider; -import com.intellij.psi.util.CachedValuesManager; -import com.intellij.psi.util.PsiModificationTracker; -import com.intellij.psi.util.PsiTreeUtil; -import com.intellij.psi.util.QualifiedName; -import com.intellij.util.ArrayUtil; -import com.intellij.util.ObjectUtils; -import com.intellij.util.containers.ContainerUtil; -import com.intellij.util.containers.Stack; -import com.jetbrains.python.PyCustomType; -import com.jetbrains.python.PyNames; -import com.jetbrains.python.PyTokenTypes; -import com.jetbrains.python.ast.PyAstFunction; -import com.jetbrains.python.ast.PyAstTypeParameter; -import com.jetbrains.python.ast.impl.PyUtilCore; -import com.jetbrains.python.codeInsight.controlflow.ScopeOwner; -import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil; -import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyFunctionTypeAnnotation; -import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyFunctionTypeAnnotationFile; -import com.jetbrains.python.codeInsight.typeHints.PyTypeHintFile; -import com.jetbrains.python.codeInsight.typeRepresentation.psi.PyFunctionTypeRepresentation; -import com.jetbrains.python.psi.AccessDirection; -import com.jetbrains.python.psi.FutureFeature; -import com.jetbrains.python.psi.LanguageLevel; -import com.jetbrains.python.psi.PyAnnotation; -import com.jetbrains.python.psi.PyAnnotationOwner; -import com.jetbrains.python.psi.PyAssignmentStatement; -import com.jetbrains.python.psi.PyBinaryExpression; -import com.jetbrains.python.psi.PyCallExpression; -import com.jetbrains.python.psi.PyCallSiteExpression; -import com.jetbrains.python.psi.PyCallable; -import com.jetbrains.python.psi.PyClass; -import com.jetbrains.python.psi.PyDecoratable; -import com.jetbrains.python.psi.PyDoubleStarExpression; -import com.jetbrains.python.psi.PyElement; -import com.jetbrains.python.psi.PyEllipsisLiteralExpression; -import com.jetbrains.python.psi.PyExpression; -import com.jetbrains.python.psi.PyExpressionCodeFragment; -import com.jetbrains.python.psi.PyFile; -import com.jetbrains.python.psi.PyForPart; -import com.jetbrains.python.psi.PyFunction; -import com.jetbrains.python.psi.PyKeywordArgument; -import com.jetbrains.python.psi.PyKnownDecorator; -import com.jetbrains.python.psi.PyKnownDecoratorUtil; -import com.jetbrains.python.psi.PyListLiteralExpression; -import com.jetbrains.python.psi.PyNamedParameter; -import com.jetbrains.python.psi.PyNoneLiteralExpression; -import com.jetbrains.python.psi.PyParameter; -import com.jetbrains.python.psi.PyParenthesizedExpression; -import com.jetbrains.python.psi.PyPsiFacade; -import com.jetbrains.python.psi.PyQualifiedNameOwner; -import com.jetbrains.python.psi.PyReferenceExpression; -import com.jetbrains.python.psi.PySequenceExpression; -import com.jetbrains.python.psi.PyStarExpression; -import com.jetbrains.python.psi.PyStatement; -import com.jetbrains.python.psi.PyStringLiteralExpression; -import com.jetbrains.python.psi.PySubscriptionExpression; -import com.jetbrains.python.psi.PyTargetExpression; -import com.jetbrains.python.psi.PyTupleExpression; -import com.jetbrains.python.psi.PyTypeAliasStatement; -import com.jetbrains.python.psi.PyTypeCommentOwner; -import com.jetbrains.python.psi.PyTypeParameter; -import com.jetbrains.python.psi.PyTypeParameterList; -import com.jetbrains.python.psi.PyTypedElement; -import com.jetbrains.python.psi.PyUtil; -import com.jetbrains.python.psi.PyWithAncestors; -import com.jetbrains.python.psi.PyWithItem; -import com.jetbrains.python.psi.impl.PyBuiltinCache; -import com.jetbrains.python.psi.impl.PyEvaluator; -import com.jetbrains.python.psi.impl.PyPsiUtils; -import com.jetbrains.python.psi.impl.stubs.PyTypingAliasStubType; -import com.jetbrains.python.psi.resolve.PyResolveContext; -import com.jetbrains.python.psi.resolve.PyResolveUtil; -import com.jetbrains.python.psi.resolve.RatedResolveResult; -import com.jetbrains.python.psi.stubs.PyClassStub; -import com.jetbrains.python.psi.stubs.PyModuleNameIndex; -import com.jetbrains.python.psi.types.PyCallableParameter; -import com.jetbrains.python.psi.types.PyCallableParameterImpl; -import com.jetbrains.python.psi.types.PyCallableParameterListType; -import com.jetbrains.python.psi.types.PyCallableParameterListTypeImpl; -import com.jetbrains.python.psi.types.PyCallableParameterVariadicType; -import com.jetbrains.python.psi.types.PyCallableTypeImpl; -import com.jetbrains.python.psi.types.PyClassLikeType; -import com.jetbrains.python.psi.types.PyClassType; -import com.jetbrains.python.psi.types.PyClassTypeImpl; -import com.jetbrains.python.psi.types.PyCollectionType; -import com.jetbrains.python.psi.types.PyCollectionTypeImpl; -import com.jetbrains.python.psi.types.PyConcatenateType; -import com.jetbrains.python.psi.types.PyInstantiableType; -import com.jetbrains.python.psi.types.PyIntersectionType; -import com.jetbrains.python.psi.types.PyLiteralStringType; -import com.jetbrains.python.psi.types.PyLiteralType; -import com.jetbrains.python.psi.types.PyModuleType; -import com.jetbrains.python.psi.types.PyNarrowedType; -import com.jetbrains.python.psi.types.PyNeverType; -import com.jetbrains.python.psi.types.PyParamSpecType; -import com.jetbrains.python.psi.types.PyPositionalVariadicType; -import com.jetbrains.python.psi.types.PySelfType; -import com.jetbrains.python.psi.types.PyTupleType; -import com.jetbrains.python.psi.types.PyType; -import com.jetbrains.python.psi.types.PyTypeChecker; -import com.jetbrains.python.psi.types.PyTypeParameterMapping; -import com.jetbrains.python.psi.types.PyTypeParameterMapping.Option; -import com.jetbrains.python.psi.types.PyTypeParameterType; -import com.jetbrains.python.psi.types.PyTypeParser; -import com.jetbrains.python.psi.types.PyTypeUtil; -import com.jetbrains.python.psi.types.PyTypeVarTupleType; -import com.jetbrains.python.psi.types.PyTypeVarTupleTypeImpl; -import com.jetbrains.python.psi.types.PyTypeVarType; -import com.jetbrains.python.psi.types.PyTypeVarTypeImpl; -import com.jetbrains.python.psi.types.PyTypedDictType; -import com.jetbrains.python.psi.types.PyTypingNewType; -import com.jetbrains.python.psi.types.PyUnionType; -import com.jetbrains.python.psi.types.PyUnpackedTupleTypeImpl; -import com.jetbrains.python.psi.types.PyVariadicType; -import com.jetbrains.python.psi.types.TypeEvalContext; -import com.jetbrains.python.sdk.legacy.PythonSdkUtil; -import one.util.streamex.StreamEx; -import org.jetbrains.annotations.ApiStatus; -import org.jetbrains.annotations.NotNull; -import org.jetbrains.annotations.Nullable; +import com.dynatrace.hash4j.hashing.HashValue128 +import com.dynatrace.hash4j.hashing.Hashing +import com.google.common.collect.ImmutableMap +import com.google.common.collect.ImmutableSet +import com.intellij.openapi.application.PathManager +import com.intellij.openapi.util.Computable +import com.intellij.openapi.util.Key +import com.intellij.openapi.util.NlsSafe +import com.intellij.openapi.util.RecursionManager +import com.intellij.openapi.util.Ref +import com.intellij.openapi.util.TextRange +import com.intellij.psi.PsiComment +import com.intellij.psi.PsiElement +import com.intellij.psi.impl.source.resolve.FileContextUtil +import com.intellij.psi.search.GlobalSearchScope +import com.intellij.psi.util.CachedValueProvider +import com.intellij.psi.util.CachedValuesManager +import com.intellij.psi.util.PsiModificationTracker +import com.intellij.psi.util.PsiTreeUtil +import com.intellij.util.ArrayUtil +import com.intellij.util.Function +import com.intellij.util.containers.ContainerUtil +import com.intellij.util.containers.Stack +import com.jetbrains.python.PyCustomType +import com.jetbrains.python.PyNames +import com.jetbrains.python.PyTokenTypes +import com.jetbrains.python.ast.PyAstFunction +import com.jetbrains.python.ast.PyAstTypeParameter +import com.jetbrains.python.ast.impl.PyUtilCore +import com.jetbrains.python.codeInsight.controlflow.ScopeOwner +import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil +import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyFunctionTypeAnnotation +import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyFunctionTypeAnnotationFile +import com.jetbrains.python.codeInsight.typeHints.PyTypeHintFile +import com.jetbrains.python.codeInsight.typeRepresentation.PyModuleTypeName +import com.jetbrains.python.codeInsight.typeRepresentation.psi.PyFunctionTypeRepresentation +import com.jetbrains.python.codeInsight.typing.PyTypeHintProvider.Companion.parseTypeHint +import com.jetbrains.python.codeInsight.typing.PyTypedDictTypeProvider.Companion.getTypedDictTypeForResolvedElement +import com.jetbrains.python.codeInsight.typing.PyTypedDictTypeProvider.Companion.isTypedDict +import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.Context +import com.jetbrains.python.psi.AccessDirection +import com.jetbrains.python.psi.FutureFeature +import com.jetbrains.python.psi.LanguageLevel +import com.jetbrains.python.psi.PyAnnotation +import com.jetbrains.python.psi.PyAnnotationOwner +import com.jetbrains.python.psi.PyAssignmentStatement +import com.jetbrains.python.psi.PyBinaryExpression +import com.jetbrains.python.psi.PyCallExpression +import com.jetbrains.python.psi.PyCallSiteExpression +import com.jetbrains.python.psi.PyCallable +import com.jetbrains.python.psi.PyClass +import com.jetbrains.python.psi.PyDecoratable +import com.jetbrains.python.psi.PyDoubleStarExpression +import com.jetbrains.python.psi.PyEllipsisLiteralExpression +import com.jetbrains.python.psi.PyExpression +import com.jetbrains.python.psi.PyFile +import com.jetbrains.python.psi.PyForPart +import com.jetbrains.python.psi.PyFunction +import com.jetbrains.python.psi.PyKeywordArgument +import com.jetbrains.python.psi.PyKnownDecorator +import com.jetbrains.python.psi.PyKnownDecoratorUtil +import com.jetbrains.python.psi.PyListLiteralExpression +import com.jetbrains.python.psi.PyNamedParameter +import com.jetbrains.python.psi.PyNoneLiteralExpression +import com.jetbrains.python.psi.PyParameter +import com.jetbrains.python.psi.PyParenthesizedExpression +import com.jetbrains.python.psi.PyPsiFacade +import com.jetbrains.python.psi.PyQualifiedNameOwner +import com.jetbrains.python.psi.PyReferenceExpression +import com.jetbrains.python.psi.PySequenceExpression +import com.jetbrains.python.psi.PyStarExpression +import com.jetbrains.python.psi.PyStatement +import com.jetbrains.python.psi.PyStringLiteralExpression +import com.jetbrains.python.psi.PySubscriptionExpression +import com.jetbrains.python.psi.PyTargetExpression +import com.jetbrains.python.psi.PyTupleExpression +import com.jetbrains.python.psi.PyTypeAliasStatement +import com.jetbrains.python.psi.PyTypeCommentOwner +import com.jetbrains.python.psi.PyTypeParameter +import com.jetbrains.python.psi.PyTypedElement +import com.jetbrains.python.psi.PyUtil +import com.jetbrains.python.psi.PyWithAncestors +import com.jetbrains.python.psi.PyWithItem +import com.jetbrains.python.psi.impl.PyBuiltinCache.Companion.getInstance +import com.jetbrains.python.psi.impl.PyEvaluator +import com.jetbrains.python.psi.impl.PyPsiUtils +import com.jetbrains.python.psi.impl.stubs.PyTypingAliasStubType +import com.jetbrains.python.psi.resolve.PyResolveContext +import com.jetbrains.python.psi.resolve.PyResolveUtil +import com.jetbrains.python.psi.resolve.RatedResolveResult +import com.jetbrains.python.psi.stubs.PyModuleNameIndex +import com.jetbrains.python.psi.types.PyCallableParameterImpl +import com.jetbrains.python.psi.types.PyCallableParameterListType +import com.jetbrains.python.psi.types.PyCallableParameterListTypeImpl +import com.jetbrains.python.psi.types.PyCallableParameterVariadicType +import com.jetbrains.python.psi.types.PyCallableTypeImpl +import com.jetbrains.python.psi.types.PyClassLikeType +import com.jetbrains.python.psi.types.PyClassType +import com.jetbrains.python.psi.types.PyClassTypeImpl +import com.jetbrains.python.psi.types.PyCollectionType +import com.jetbrains.python.psi.types.PyCollectionTypeImpl +import com.jetbrains.python.psi.types.PyConcatenateType +import com.jetbrains.python.psi.types.PyInstantiableType +import com.jetbrains.python.psi.types.PyIntersectionType.Companion.intersection +import com.jetbrains.python.psi.types.PyLiteralStringType.Companion.create +import com.jetbrains.python.psi.types.PyLiteralType +import com.jetbrains.python.psi.types.PyModuleType +import com.jetbrains.python.psi.types.PyNarrowedType +import com.jetbrains.python.psi.types.PyNarrowedType.Companion.create +import com.jetbrains.python.psi.types.PyNeverType +import com.jetbrains.python.psi.types.PyParamSpecType +import com.jetbrains.python.psi.types.PyPositionalVariadicType +import com.jetbrains.python.psi.types.PySelfType +import com.jetbrains.python.psi.types.PyTupleType +import com.jetbrains.python.psi.types.PyType +import com.jetbrains.python.psi.types.PyTypeChecker +import com.jetbrains.python.psi.types.PyTypeParameterMapping +import com.jetbrains.python.psi.types.PyTypeParameterType +import com.jetbrains.python.psi.types.PyTypeParser +import com.jetbrains.python.psi.types.PyTypeUtil +import com.jetbrains.python.psi.types.PyTypeVarTupleType +import com.jetbrains.python.psi.types.PyTypeVarTupleTypeImpl +import com.jetbrains.python.psi.types.PyTypeVarType +import com.jetbrains.python.psi.types.PyTypeVarTypeImpl +import com.jetbrains.python.psi.types.PyTypedDictType +import com.jetbrains.python.psi.types.PyUnionType +import com.jetbrains.python.psi.types.PyUnpackedTupleTypeImpl +import com.jetbrains.python.psi.types.PyVariadicType +import com.jetbrains.python.psi.types.TypeEvalContext +import com.jetbrains.python.sdk.legacy.PythonSdkUtil +import one.util.streamex.StreamEx +import org.jetbrains.annotations.ApiStatus +import java.nio.file.Path +import java.util.Collections +import java.util.Objects +import java.util.Optional +import java.util.function.UnaryOperator +import java.util.regex.Pattern -import java.nio.file.Path; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collection; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.LinkedHashMap; -import java.util.LinkedHashSet; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Optional; -import java.util.Set; -import java.util.function.Function; -import java.util.regex.Pattern; -import java.util.stream.Stream; - -import static com.dynatrace.hash4j.hashing.Hashing.xxh3_128; -import static com.intellij.openapi.util.RecursionManager.doPreventingRecursion; -import static com.jetbrains.python.codeInsight.typeRepresentation.PyTypeRepresentationDialectKt.PyModuleTypeName; -import static com.jetbrains.python.psi.PyKnownDecorator.TYPING_FINAL; -import static com.jetbrains.python.psi.PyKnownDecorator.TYPING_FINAL_EXT; -import static com.jetbrains.python.psi.PyUtil.as; - -public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext { - - public static final @NlsSafe String TYPING = "typing"; - - public static final String GENERATOR = "typing.Generator"; - public static final String ASYNC_GENERATOR = "typing.AsyncGenerator"; - public static final String COROUTINE = "typing.Coroutine"; - public static final String AWAITABLE = "typing.Awaitable"; - public static final String NAMEDTUPLE = "typing.NamedTuple"; - public static final String TYPED_DICT = "typing.TypedDict"; - public static final String TYPED_DICT_EXT = "typing_extensions.TypedDict"; - public static final String TYPE_GUARD = "typing.TypeGuard"; - public static final String TYPE_GUARD_EXT = "typing_extensions.TypeGuard"; - public static final String TYPE_IS = "typing.TypeIs"; - public static final String TYPE_IS_EXT = "typing_extensions.TypeIs"; - public static final String GENERIC = "typing.Generic"; - public static final String PROTOCOL = "typing.Protocol"; - public static final String PROTOCOL_EXT = "typing_extensions.Protocol"; - public static final String TYPE = "typing.Type"; - public static final String ANY = "typing.Any"; - public static final String NEW_TYPE = "typing.NewType"; - public static final String CALLABLE = "typing.Callable"; - public static final String CALLABLE_EXT = "typing_extensions.Callable"; - public static final String MAPPING = "typing.Mapping"; - public static final String MAPPING_GET = "typing.Mapping.get"; - private static final String LIST = "typing.List"; - private static final String DICT = "typing.Dict"; - private static final String DEFAULT_DICT = "typing.DefaultDict"; - private static final String ORDERED_DICT = "typing.OrderedDict"; - private static final String SET = "typing.Set"; - private static final String FROZEN_SET = "typing.FrozenSet"; - private static final String COUNTER = "typing.Counter"; - private static final String DEQUE = "typing.Deque"; - private static final String TUPLE = "typing.Tuple"; - public static final String CLASS_VAR = "typing.ClassVar"; - public static final String TYPE_VAR = "typing.TypeVar"; - public static final String TYPE_VAR_EXT = "typing_extensions.TypeVar"; - public static final String TYPE_VAR_TUPLE = "typing.TypeVarTuple"; - public static final String TYPE_VAR_TUPLE_EXT = "typing_extensions.TypeVarTuple"; - public static final String PARAM_SPEC = "typing.ParamSpec"; - public static final String PARAM_SPEC_EXT = "typing_extensions.ParamSpec"; - private static final String CHAIN_MAP = "typing.ChainMap"; - public static final String UNION = "typing.Union"; - public static final String CONCATENATE = "typing.Concatenate"; - public static final String CONCATENATE_EXT = "typing_extensions.Concatenate"; - public static final String OPTIONAL = "typing.Optional"; - public static final String NO_RETURN = "typing.NoReturn"; - public static final String NEVER = "typing.Never"; - public static final String NO_RETURN_EXT = "typing_extensions.NoReturn"; - public static final String NEVER_EXT = "typing_extensions.Never"; - public static final String FINAL = "typing.Final"; - public static final String FINAL_EXT = "typing_extensions.Final"; - public static final String LITERAL = "typing.Literal"; - public static final String LITERAL_EXT = "typing_extensions.Literal"; - public static final String LITERALSTRING = "typing.LiteralString"; - public static final String LITERALSTRING_EXT = "typing_extensions.LiteralString"; - public static final String ANNOTATED = "typing.Annotated"; - public static final String ANNOTATED_EXT = "typing_extensions.Annotated"; - public static final String TYPE_ALIAS = "typing.TypeAlias"; - public static final String TYPE_ALIAS_EXT = "typing_extensions.TypeAlias"; - public static final String TYPE_ALIAS_TYPE = "typing.TypeAliasType"; - public static final String SPECIAL_FORM = "typing._SpecialForm"; - public static final String SPECIAL_FORM_EXT = "typing_extensions._SpecialForm"; - public static final String REQUIRED = "typing.Required"; - public static final String REQUIRED_EXT = "typing_extensions.Required"; - public static final String NOT_REQUIRED = "typing.NotRequired"; - public static final String NOT_REQUIRED_EXT = "typing_extensions.NotRequired"; - public static final String READONLY = "typing.ReadOnly"; - public static final String READONLY_EXT = "typing_extensions.ReadOnly"; - - public static final Set TYPE_PARAMETER_FACTORIES = Set.of( - TYPE_VAR, TYPE_VAR_EXT, - PARAM_SPEC, PARAM_SPEC_EXT, - TYPE_VAR_TUPLE, TYPE_VAR_TUPLE_EXT - ); - - public static final Set TYPE_DICT_QUALIFIERS = - Set.of(REQUIRED, REQUIRED_EXT, NOT_REQUIRED, NOT_REQUIRED_EXT, READONLY, READONLY_EXT); - - public static final String UNPACK = "typing.Unpack"; - public static final String UNPACK_EXT = "typing_extensions.Unpack"; - - public static final String SELF = "typing.Self"; - public static final String SELF_EXT = "typing_extensions.Self"; - - public static final Pattern TYPE_IGNORE_PATTERN = - Pattern.compile("#\\s*type:\\s*ignore\\s*(\\[[^]#]*])?($|(\\s.*))", Pattern.CASE_INSENSITIVE); - - public static final String ASSERT_TYPE = "typing.assert_type"; - public static final String REVEAL_TYPE = "typing.reveal_type"; - public static final String REVEAL_TYPE_EXT = "typing_extensions.reveal_type"; - public static final String CAST = "typing.cast"; - public static final String CAST_EXT = "typing_extensions.cast"; - - public static final ImmutableMap BUILTIN_COLLECTION_CLASSES = ImmutableMap.builder() - .put(LIST, "list") - .put(DICT, "dict") - .put(SET, PyNames.SET) - .put(FROZEN_SET, "frozenset") - .put(TUPLE, PyNames.TUPLE) - .build(); - - private static final ImmutableMap COLLECTIONS_CLASSES = ImmutableMap.builder() - .put(DEFAULT_DICT, "collections.defaultdict") - .put(ORDERED_DICT, "collections.OrderedDict") - .put(COUNTER, "collections.Counter") - .put(DEQUE, "collections.deque") - .put(CHAIN_MAP, "collections.ChainMap") - .build(); - - public static final ImmutableMap TYPING_COLLECTION_CLASSES = ImmutableMap.builder() - .put("list", "List") - .put("dict", "Dict") - .put("set", "Set") - .put("frozenset", "FrozenSet") - .build(); - - public static final ImmutableMap TYPING_BUILTINS_GENERIC_ALIASES = ImmutableMap.builder() - .putAll(TYPING_COLLECTION_CLASSES.entrySet()) - .put("type", "Type") - .put("tuple", "Tuple") - .build(); - - public static final ImmutableSet GENERIC_CLASSES = ImmutableSet.builder() - // special forms - .add(TUPLE, GENERIC, PROTOCOL, CALLABLE, CALLABLE_EXT, TYPE, CLASS_VAR, FINAL, LITERAL, ANNOTATED, REQUIRED, NOT_REQUIRED, READONLY) - // type aliases - .add(UNION, OPTIONAL, LIST, DICT, DEFAULT_DICT, ORDERED_DICT, SET, FROZEN_SET, COUNTER, DEQUE, CHAIN_MAP) - .add(PROTOCOL_EXT, FINAL_EXT, LITERAL_EXT, ANNOTATED_EXT, REQUIRED_EXT, NOT_REQUIRED_EXT, READONLY_EXT) - .build(); - - /** - * For the following names we shouldn't go further to the RHS of assignments, - * since they are not type aliases already and in typing.pyi are assigned to - * some synthetic values. - */ - public static final ImmutableSet OPAQUE_NAMES = ImmutableSet.builder() - .add(PyKnownDecorator.TYPING_OVERLOAD.getQualifiedName().toString()) - .add(ANY) - .add(TYPE_VAR) - .add(TYPE_VAR_EXT) - .add(TYPE_VAR_TUPLE) - .add(TYPE_VAR_TUPLE_EXT) - .add(GENERIC) - .add(PARAM_SPEC) - .add(PARAM_SPEC_EXT) - .add(CONCATENATE) - .add(CONCATENATE_EXT) - .add(TUPLE) - .add(CALLABLE) - .add(CALLABLE_EXT) - .add(TYPE) - .add(PyKnownDecorator.TYPING_NO_TYPE_CHECK.getQualifiedName().toString()) - .add(PyKnownDecorator.TYPING_NO_TYPE_CHECK_EXT.getQualifiedName().toString()) - .add(UNION) - .add(OPTIONAL) - .add(LIST) - .add(DICT) - .add(DEFAULT_DICT) - .add(ORDERED_DICT) - .add(SET) - .add(FROZEN_SET) - .add(PROTOCOL, PROTOCOL_EXT) - .add(CLASS_VAR) - .add(COUNTER) - .add(DEQUE) - .add(CHAIN_MAP) - .add(NO_RETURN, NO_RETURN_EXT) - .add(NEVER, NEVER_EXT) - .add(FINAL, FINAL_EXT) - .add(LITERAL, LITERAL_EXT) - .add(TYPED_DICT, TYPED_DICT_EXT) - .add(ANNOTATED, ANNOTATED_EXT) - .add(TYPE_ALIAS, TYPE_ALIAS_EXT) - .add(REQUIRED, REQUIRED_EXT) - .add(NOT_REQUIRED, NOT_REQUIRED_EXT) - .add(READONLY, READONLY_EXT) - .add(SELF, SELF_EXT) - .add(LITERALSTRING, LITERALSTRING_EXT) - .build(); - - private static final Key TYPE_HINT_EVAL_CONTEXT = Key.create("TYPE_HINT_EVAL_CONTEXT"); - - @Override - public @Nullable PyType getReferenceExpressionType(@NotNull PyReferenceExpression referenceExpression, @NotNull Context context) { +class PyTypingTypeProvider : PyTypeProviderWithCustomContext() { + public override fun getReferenceExpressionType(referenceExpression: PyReferenceExpression, context: Context): PyType? { // Check for the exact name in advance for performance reasons - if ("Generic".equals(referenceExpression.getName())) { - if (resolvesToQualifiedNames(referenceExpression, context.myContext, GENERIC)) { - return createTypingGenericType(referenceExpression); + if ("Generic" == referenceExpression.name) { + if (resolvesToQualifiedNames(referenceExpression, context.typeContext, GENERIC)) { + return createTypingGenericType(referenceExpression) } } // Check for the exact name in advance for performance reasons - if ("Protocol".equals(referenceExpression.getName())) { - if (resolvesToQualifiedNames(referenceExpression, context.myContext, PROTOCOL, PROTOCOL_EXT)) { - return createTypingProtocolType(referenceExpression); + if ("Protocol" == referenceExpression.name) { + if (resolvesToQualifiedNames(referenceExpression, context.typeContext, PROTOCOL, PROTOCOL_EXT)) { + return createTypingProtocolType(referenceExpression) } } // Check for the exact name in advance for performance reasons - if ("Callable".equals(referenceExpression.getName())) { - if (resolvesToQualifiedNames(referenceExpression, context.myContext, CALLABLE, CALLABLE_EXT)) { - return createTypingCallableType(referenceExpression); + if ("Callable" == referenceExpression.name) { + if (resolvesToQualifiedNames(referenceExpression, context.typeContext, CALLABLE, CALLABLE_EXT)) { + return createTypingCallableType(referenceExpression) } } - return null; + return null } - @Override - public @Nullable Ref getParameterType(@NotNull PyNamedParameter param, @NotNull PyFunction func, @NotNull Context context) { - @Nullable PyExpression typeHint = getAnnotationValue(param, context.myContext); + override fun getParameterType(param: PyNamedParameter, func: PyFunction, context: Context): Ref? { + var typeHint: PyExpression? = getAnnotationValue(param, context.typeContext) if (typeHint == null) { - String paramTypeCommentHint = param.getTypeCommentAnnotation(); + val paramTypeCommentHint = param.typeCommentAnnotation if (paramTypeCommentHint != null) { - typeHint = PyUtil.createExpressionFromFragment(paramTypeCommentHint, param); + typeHint = PyUtil.createExpressionFromFragment(paramTypeCommentHint, param) } } if (typeHint == null) { - PyFunctionTypeAnnotation annotation = getFunctionTypeAnnotation(func); + val annotation: PyFunctionTypeAnnotation? = getFunctionTypeAnnotation(func) if (annotation != null) { - PyExpression funcTypeCommentParamHint = findParamTypeHintInFunctionTypeComment(annotation, param, func); + val funcTypeCommentParamHint: PyExpression? = findParamTypeHintInFunctionTypeComment(annotation, param, func) if (funcTypeCommentParamHint == null) { - return Ref.create(); + return Ref() } - typeHint = funcTypeCommentParamHint; + typeHint = funcTypeCommentParamHint } } if (typeHint == null) { - return null; + return null } - if (param.isKeywordContainer()) { - Ref type = getTypeFromUnpackOperator(typeHint, context.myContext); + if (param.isKeywordContainer) { + val type: Ref? = getTypeFromUnpackOperator(typeHint, context.typeContext) if (type != null) { - return type.get() instanceof PyTypedDictType ? type : null; + return if (type.get() is PyTypedDictType) type else null } } - if (typeHint instanceof PyReferenceExpression ref && ref.isQualified() && ( - param.isPositionalContainer() && "args".equals(ref.getReferencedName()) || - param.isKeywordContainer() && "kwargs".equals(ref.getReferencedName()) - )) { - typeHint = Objects.requireNonNull(ref.getQualifier()); + if (typeHint is PyReferenceExpression && typeHint.isQualified && (param.isPositionalContainer && "args" == typeHint.referencedName || + param.isKeywordContainer && "kwargs" == typeHint.referencedName + ) + ) { + typeHint = typeHint.qualifier!! } - PyType type = Ref.deref(getType(typeHint, context)); - if (param.isPositionalContainer() && !(type instanceof PyParamSpecType)) { - return Ref.create(PyTypeUtil.toPositionalContainerType(param, type)); + val type = Ref.deref(getType(typeHint, context)) + if (param.isPositionalContainer && type !is PyParamSpecType) { + return Ref(PyTypeUtil.toPositionalContainerType(param, type)) } - if (param.isKeywordContainer() && !(type instanceof PyParamSpecType)) { - return Ref.create(PyTypeUtil.toKeywordContainerType(param, type)); + if (param.isKeywordContainer && type !is PyParamSpecType) { + return Ref(PyTypeUtil.toKeywordContainerType(param, type)) } - if (PyNames.NONE.equals(param.getDefaultValueText())) { - return Ref.create(PyUnionType.union(type, PyBuiltinCache.getInstance(param).getNoneType())); + if (PyNames.NONE == param.defaultValueText) { + return Ref(PyUnionType.union(type, getInstance(param).noneType)) } - return Ref.create(type); + return Ref(type) } - private static @Nullable PyExpression findParamTypeHintInFunctionTypeComment(@NotNull PyFunctionTypeAnnotation annotation, - @NotNull PyNamedParameter param, - @NotNull PyFunction func) { - List paramTypes = annotation.getParameterTypeList().getParameterTypes(); - if (paramTypes.size() == 1 && paramTypes.get(0) instanceof PyEllipsisLiteralExpression) { - return null; - } - int startOffset = omitFirstParamInTypeComment(func, annotation) ? 1 : 0; - List funcParams = Arrays.asList(func.getParameterList().getParameters()); - int i = funcParams.indexOf(param) - startOffset; - if (i >= 0 && i < paramTypes.size()) { - PyExpression paramTypeHint = paramTypes.get(i); - if (paramTypeHint instanceof PyStarExpression starExpression) { - return starExpression.getExpression(); - } - else if (paramTypeHint instanceof PyDoubleStarExpression doubleStarExpression) { - return doubleStarExpression.getExpression(); - } - else { - return paramTypeHint; - } - } - return null; - } - - public static boolean isGenerator(@NotNull PyType type) { - return type instanceof PyCollectionType genericType && GENERATOR.equals(genericType.getClassQName()); - } - - private static @NotNull PyType createTypingGenericType(@NotNull PsiElement anchor) { - return new PyCustomType(GENERIC, null, false, true, PyBuiltinCache.getInstance(anchor).getObjectType()); - } - - private static @NotNull PyType createTypingProtocolType(@NotNull PsiElement anchor) { - return new PyCustomType(PROTOCOL, null, false, true, PyBuiltinCache.getInstance(anchor).getObjectType()); - } - - public static @NotNull PyType createTypingCallableType(@NotNull PsiElement anchor) { - return new PyCustomType(CALLABLE, null, false, true, PyBuiltinCache.getInstance(anchor).getObjectType()); - } - - private static boolean omitFirstParamInTypeComment(@NotNull PyFunction func, @NotNull PyFunctionTypeAnnotation annotation) { - return func.getContainingClass() != null && func.getModifier() != PyAstFunction.Modifier.STATICMETHOD && - annotation.getParameterTypeList().getParameterTypes().size() < func.getParameterList().getParameters().length; - } - - @Override - public @Nullable Ref getReturnType(@NotNull PyCallable callable, @NotNull Context context) { - if (callable instanceof PyFunction function) { - final PyExpression returnTypeAnnotation = getReturnTypeAnnotation(function, context.myContext); + override fun getReturnType(callable: PyCallable, context: Context): Ref? { + if (callable is PyFunction) { + val returnTypeAnnotation: PyExpression? = getReturnTypeAnnotation(callable, context.typeContext) if (returnTypeAnnotation != null) { - final Ref typeRef = getType(returnTypeAnnotation, context); + val typeRef: Ref? = getType(returnTypeAnnotation, context) if (typeRef != null) { // Do not use toAsyncIfNeeded, as it also converts Generators. Here we do not need it. - if (function.isAsync() && function.isAsyncAllowed() && !function.isGenerator()) { - return Ref.create(wrapInCoroutineType(typeRef.get(), function)); + if (callable.isAsync && callable.isAsyncAllowed && !callable.isGenerator) { + return Ref(wrapInCoroutineType(typeRef.get(), callable)) } - return typeRef; + return typeRef } // Don't rely on other type providers if a type hint is present, but cannot be resolved. - return Ref.create(); + return Ref() } } - return null; + return null } - @ApiStatus.Internal - public static @Nullable PyExpression getReturnTypeAnnotation(@NotNull PyFunction function, TypeEvalContext context) { - final PyExpression returnAnnotation = getAnnotationValue(function, context); - if (returnAnnotation != null) { - return returnAnnotation; - } - final PyFunctionTypeAnnotation functionAnnotation = getFunctionTypeAnnotation(function); - if (functionAnnotation != null) { - return functionAnnotation.getReturnType(); - } - return null; - } + override fun getCallType(function: PyFunction, callSite: PyCallSiteExpression, context: Context): Ref? { + val functionQName = function.qualifiedName - public static @Nullable PyFunctionTypeAnnotation getFunctionTypeAnnotation(@NotNull PyFunction function) { - final String comment = function.getTypeCommentAnnotation(); - if (comment == null) { - return null; - } - final PyFunctionTypeAnnotationFile file = CachedValuesManager.getCachedValue(function, () -> - CachedValueProvider.Result.create(new PyFunctionTypeAnnotationFile(function.getTypeCommentAnnotation(), function), function)); - return file.getAnnotation(); - } - - @Override - public @Nullable Ref getCallType(@NotNull PyFunction function, @NotNull PyCallSiteExpression callSite, @NotNull Context context) { - final String functionQName = function.getQualifiedName(); - - if (CAST.equals(functionQName) || CAST_EXT.equals(functionQName)) { + if (CAST == functionQName || CAST_EXT == functionQName) { return Optional - .ofNullable(as(callSite, PyCallExpression.class)) - .map(PyCallExpression::getArguments) - .filter(args -> args.length > 0) - .map(args -> getType(args[0], context)) - .orElse(null); + .ofNullable(callSite as PyCallExpression) + .map { it.arguments } + .filter { it.isNotEmpty() } + .map { getType(it[0]!!, context) } + .orElse(null) } if (functionReturningCallSiteAsAType(function)) { - return getAsClassObjectType(callSite, context); + return getAsClassObjectType(callSite, context) } - return null; + return null } - private static boolean functionReturningCallSiteAsAType(@NotNull PyFunction function) { - final String name = function.getName(); - - if (PyNames.CLASS_GETITEM.equals(name)) return true; - if (PyNames.GETITEM.equals(name)) { - final PyClass cls = function.getContainingClass(); - if (cls != null) { - final String qualifiedName = cls.getQualifiedName(); - return SPECIAL_FORM.equals(qualifiedName) || SPECIAL_FORM_EXT.equals(qualifiedName); - } - } - - return false; - } - - private static @Nullable PyType getTypedDictTypeForTarget(@NotNull PyTargetExpression referenceTarget, @NotNull TypeEvalContext context) { - if (PyTypedDictTypeProvider.Companion.isTypedDict(referenceTarget, context)) { - return new PyCustomType(TYPED_DICT, null, false, true, - PyBuiltinCache.getInstance(referenceTarget).getDictType()); - } - - return null; - } - - @Override - public Ref getReferenceType(@NotNull PsiElement referenceTarget, @NotNull Context context, @Nullable PsiElement anchor) { - if (referenceTarget instanceof PyTargetExpression target) { - final String targetQName = target.getQualifiedName(); + override fun getReferenceType(referenceTarget: PsiElement, context: Context, anchor: PsiElement?): Ref? { + var referenceTarget = referenceTarget + if (referenceTarget is PyTargetExpression) { + val targetQName = referenceTarget.qualifiedName // Depends on typing.Generic defined as a target expression - if (GENERIC.equals(targetQName)) { - return Ref.create(createTypingGenericType(target)); + if (GENERIC == targetQName) { + return Ref(createTypingGenericType(referenceTarget)) } // Depends on typing.Protocol defined as a target expression - if (PROTOCOL.equals(targetQName) || PROTOCOL_EXT.equals(targetQName)) { - return Ref.create(createTypingProtocolType(target)); + if (PROTOCOL == targetQName || PROTOCOL_EXT == targetQName) { + return Ref(createTypingProtocolType(referenceTarget)) } // Depends on typing.Callable defined as a target expression - if (CALLABLE.equals(targetQName) || CALLABLE_EXT.equals(targetQName)) { - return Ref.create(createTypingCallableType(referenceTarget)); + if (CALLABLE == targetQName || CALLABLE_EXT == targetQName) { + return Ref(createTypingCallableType(referenceTarget)) } - final PyType collection = getCollection(target, context.myContext); - if (collection instanceof PyInstantiableType) { - return Ref.create(((PyInstantiableType)collection).toClass()); + val collection: PyType? = getCollection(referenceTarget, context.typeContext) + if (collection is PyInstantiableType<*>) { + return Ref(collection.toClass()) } - final PyType typedDictType = getTypedDictTypeForTarget(target, context.myContext); + val typedDictType: PyType? = getTypedDictTypeForTarget(referenceTarget, context.typeContext) if (typedDictType != null) { - return Ref.create(typedDictType); + return Ref(typedDictType) } - PyExpression assignedValue = PyTypingAliasStubType.getAssignedValueStubLike(target); - if (assignedValue instanceof PyCallExpression callExpression && callExpression.isCalleeText("TypeVar", "TypeVarTuple", "ParamSpec")) { - for (PsiElement element : tryResolving(Objects.requireNonNull(callExpression.getCallee()), context.myContext)) { - if (element instanceof PyClass pyClass && - TYPE_PARAMETER_FACTORIES.contains(ObjectUtils.notNull(pyClass.getQualifiedName(), "")) && - context.myContext.getType(pyClass) instanceof PyClassType pyClassType) { - return Ref.create(pyClassType.toInstance()); + val assignedValue = PyTypingAliasStubType.getAssignedValueStubLike(referenceTarget) + if (assignedValue is PyCallExpression && assignedValue.isCalleeText("TypeVar", "TypeVarTuple", "ParamSpec")) { + for (element in tryResolving(assignedValue.callee!!, context.typeContext)) { + if ( + element is PyClass && + (element.qualifiedName ?: "") in TYPE_PARAMETER_FACTORIES + ) { + (context.typeContext.getType(element) as? PyClassType)?.let { + return Ref(it.toInstance()) + } } } } // Return a type from an immediate type hint, e.g. from a syntactic annotation for - // + // // x: int = ... // // or find the "root" declaration and get a type from a type hint there, e.g. - // + // // x: int // x = ... // @@ -610,2075 +309,2610 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext< // self.attr = ... // assignments "inst.attr = ..." are not preserved in stubs anyway. See PyTargetExpressionElementType.shouldCreateStub. - if (target.isQualified() && context.myContext.maySwitchToAST(target)) { - PsiElement resolved = target.getReference(PyResolveContext.defaultContext(context.myContext)).resolve(); - if (resolved instanceof PyTargetExpression resolvedTarget && PyUtil.isAttribute(resolvedTarget)) { - target = resolvedTarget; + if (referenceTarget.isQualified && context.typeContext.maySwitchToAST(referenceTarget)) { + val resolved = referenceTarget.getReference(PyResolveContext.defaultContext(context.typeContext)).resolve() + if (resolved is PyTargetExpression && PyUtil.isAttribute(resolved)) { + referenceTarget = resolved } } - final Ref annotatedType = getTypeFromTypeHint(target, context); + val annotatedType: Ref? = getTypeFromTypeHint(referenceTarget, context) if (annotatedType != null) { - return annotatedType; + return annotatedType } - final String name = target.getReferencedName(); - final ScopeOwner scopeOwner = ScopeUtil.getScopeOwner(target); + val name = referenceTarget.referencedName + val scopeOwner = ScopeUtil.getScopeOwner(referenceTarget) if (name == null || scopeOwner == null) { - return null; + return null } - final PyClass pyClass = target.getContainingClass(); + val pyClass = referenceTarget.containingClass - if (target.isQualified()) { - if (pyClass != null && scopeOwner instanceof PyFunction) { - final PyResolveContext resolveContext = PyResolveContext.defaultContext(context.myContext); + if (referenceTarget.isQualified) { + if (pyClass != null && scopeOwner is PyFunction) { + val resolveContext = PyResolveContext.defaultContext(context.typeContext) - boolean isInstanceAttribute; - if (context.myContext.maySwitchToAST(target)) { - isInstanceAttribute = StreamEx.of(PyUtil.multiResolveTopPriority(target.getQualifier(), resolveContext)) - .select(PyParameter.class) - .filter(PyParameter::isSelf) - .anyMatch(p -> PsiTreeUtil.getParentOfType(p, PyFunction.class) == scopeOwner); + val isInstanceAttribute: Boolean + if (context.typeContext.maySwitchToAST(referenceTarget)) { + isInstanceAttribute = + StreamEx.of(PyUtil.multiResolveTopPriority(referenceTarget.qualifier!!, resolveContext)) + .select(PyParameter::class.java) + .filter { obj: PyParameter? -> obj!!.isSelf } + .anyMatch { p: PyParameter? -> + PsiTreeUtil.getParentOfType( + p, + PyFunction::class.java + ) === scopeOwner + } } else { - isInstanceAttribute = PyUtil.isInstanceAttribute(target); + isInstanceAttribute = PyUtil.isInstanceAttribute(referenceTarget) } if (!isInstanceAttribute) { - return null; + return null } // Set isDefinition=true to start searching right from the class level. - Ref memberType = - getMemberTypeForClassType(context, target, name, resolveContext, false, new PyClassTypeImpl(pyClass, true)); + val memberType: Ref? = + getMemberTypeForClassType(context, referenceTarget, name, resolveContext, false, PyClassTypeImpl(pyClass, true)) if (memberType != null) { - return memberType; + return memberType } - for (PyClass ancestor : pyClass.getAncestorClasses(resolveContext.getTypeEvalContext())) { - Ref ancestorMemberType = - getMemberTypeForClassType(context, target, name, resolveContext, true, new PyClassTypeImpl(ancestor, false)); + for (ancestor in pyClass.getAncestorClasses(resolveContext.typeEvalContext)) { + val ancestorMemberType: Ref? = + getMemberTypeForClassType( + context, + referenceTarget, + name, + resolveContext, + true, + PyClassTypeImpl(ancestor, false) + ) if (ancestorMemberType != null) { - return ancestorMemberType; + return ancestorMemberType } } - return null; + return null } } else { - if (context.myContext.maySwitchToAST(target)) { - final PyResolveContext resolveContext = PyResolveContext.defaultContext(context.myContext); - final Set visited = new HashSet<>(); - PsiElement current = target; - while (current instanceof PyTargetExpression currentTarget && visited.add(currentTarget)) { - List resolveResults = Arrays.stream(currentTarget.getReference(resolveContext).multiResolve(false)) - .filter(r -> r.getElement() != currentTarget) - .toList(); + if (context.typeContext.maySwitchToAST(referenceTarget)) { + val resolveContext = PyResolveContext.defaultContext(context.typeContext) + val visited: MutableSet = HashSet() + var current: PsiElement? = referenceTarget + while (current is PyTargetExpression && visited.add(current)) { + val resolveResults = current + .getReference(resolveContext) + .multiResolve(false) + .filter { it.element !== current } - current = ContainerUtil.getFirstItem(PyUtil.filterTopPriorityElements(resolveResults)); + current = PyUtil.filterTopPriorityElements(resolveResults).firstOrNull() - if (current instanceof PyTargetExpression) { - Ref type = getTypeFromTypeHint((PyTargetExpression)current, context); + if (current is PyTargetExpression) { + val type = getTypeFromTypeHint(current, context) if (type != null) { - return type; + return type } } - else if (current instanceof PyNamedParameter namedParameter) { - Ref type = getTypeFromTypeHint(namedParameter, context); + else if (current is PyNamedParameter) { + val type = getTypeFromTypeHint(current, context) if (type != null) { - if (namedParameter.isPositionalContainer()) { - return Ref.create(PyTypeUtil.toPositionalContainerType(namedParameter, type.get())); - } else if (namedParameter.isKeywordContainer()) { - return Ref.create(PyTypeUtil.toKeywordContainerType(namedParameter, type.get())); + if (current.isPositionalContainer) { + return Ref.create(PyTypeUtil.toPositionalContainerType(current, type.get())) } - return type; + else if (current.isKeywordContainer) { + return Ref.create(PyTypeUtil.toKeywordContainerType(current, type.get())) + } + return type } } } } else { - List candidates = null; - if (scopeOwner instanceof PyFile) { - candidates = ((PyFile)scopeOwner).getTopLevelAttributes(); - } - else if (scopeOwner instanceof PyClass) { - candidates = ((PyClass)scopeOwner).getClassAttributes(); - } + val candidates = + when (scopeOwner) { + is PyFile -> scopeOwner.topLevelAttributes + is PyClass -> scopeOwner.classAttributes + else -> null + } if (candidates != null) { - return StreamEx.of(candidates) - .filter(t -> name.equals(t.getName())) - .map(x -> getTypeFromTypeHint(x, context)) - .nonNull() - .findFirst() - .orElse(null); + return candidates + .filter { name == it.name } + .firstNotNullOfOrNull { getTypeFromTypeHint(it, context) } } } } } - if (anchor instanceof PyExpression && referenceTarget instanceof PyClass pyClass && isGeneric(pyClass, context.myContext)) { - PyCollectionType parameterizedType = parameterizeClassDefaultAware(pyClass, List.of(), context); + if (anchor is PyExpression && referenceTarget is PyClass && isGeneric(referenceTarget, context.typeContext)) { + val parameterizedType = parameterizeClassDefaultAware(referenceTarget, listOf(), context) if (parameterizedType != null) { - return Ref.create(parameterizedType.toClass()); + return Ref(parameterizedType.toClass()) } } - return null; + return null } - private static @Nullable Ref getMemberTypeForClassType(@NotNull Context context, - PyTargetExpression target, - String name, - PyResolveContext resolveContext, - boolean isInherited, - PyClassTypeImpl classType) { - final List classAttrs = - classType.resolveMember(name, target, AccessDirection.READ, resolveContext, isInherited); - if (classAttrs == null) { - return null; + override fun getGenericType(cls: PyClass, context: Context): PyType? { + val typeParameters = collectTypeParameters(cls, context) + return if (typeParameters.isEmpty()) null else PyCollectionTypeImpl(cls, false, typeParameters) + } + + override fun getGenericSubstitutions(cls: PyClass, context: Context): Map { + return PyUtil.getParameterizedCachedValue( + cls, + context + ) { calculateGenericSubstitutions(cls, it) } + } + + private fun calculateGenericSubstitutions(cls: PyClass, context: Context): Map { + if (!isGeneric(cls, context.typeContext)) { + return emptyMap() } - return StreamEx.of(classAttrs) - .map(RatedResolveResult::getElement) - .select(PyTargetExpression.class) - .filter(x -> { - ScopeOwner owner = ScopeUtil.getScopeOwner(x); - return owner instanceof PyClass || owner instanceof PyFunction; - }) - .map(x -> getTypeFromTypeHint(x, context)) - .collect(PyTypeUtil.toUnionFromRef()); - } - - private static - @Nullable Ref getTypeFromTypeHint(@NotNull T element, @NotNull Context context) { - final PyExpression annotation = getAnnotationValue(element, context.myContext); - if (annotation != null) { - return getType(annotation, context); - } - final String comment = element.getTypeCommentAnnotation(); - if (comment != null) { - return getVariableTypeCommentType(comment, element, context); - } - return null; - } - - /** - * Checks that text of a comment starts with "# type:" prefix and returns trimmed type hint after it. - * The trailing part is supposed to contain type annotation in PEP 484 compatible format and an optional - * plain text comment separated from it with another "#". - *

- * For instance, for {@code # type: List[int] # comment} it returns {@code List[int]}. - *

- * This method cannot return an empty string. - * - * @see #getTypeCommentValueRange(String) - */ - public static @Nullable String getTypeCommentValue(@NotNull String text) { - return PyUtilCore.getTypeCommentValue(text); - } - - /** - * Returns the corresponding text range for a type hint as returned by {@link #getTypeCommentValue(String)}. - * - * @see #getTypeCommentValue(String) - */ - public static @Nullable TextRange getTypeCommentValueRange(@NotNull String text) { - return PyUtilCore.getTypeCommentValueRange(text); - } - - @Override - public @Nullable PyType getGenericType(@NotNull PyClass cls, @NotNull Context context) { - List typeParameters = collectTypeParameters(cls, context); - return typeParameters.isEmpty() ? null : new PyCollectionTypeImpl(cls, false, typeParameters); - } - - @Override - public @NotNull Map getGenericSubstitutions(@NotNull PyClass cls, @NotNull Context context) { - return PyUtil.getParameterizedCachedValue(cls, context, c -> calculateGenericSubstitutions(cls, c)); - } - - private @NotNull Map calculateGenericSubstitutions(@NotNull PyClass cls, @NotNull Context context) { - if (!isGeneric(cls, context.myContext)) { - return Collections.emptyMap(); - } - Map results = new HashMap<>(); - for (PyClassType superClassType : evaluateSuperClassesAsTypeHints(cls, context.myContext)) { - Map superSubstitutions = - doPreventingRecursion(superClassType.getPyClass(), false, () -> getGenericSubstitutions(superClassType.getPyClass(), context)); - if (superSubstitutions != null) { - results.putAll(superSubstitutions); + val results = HashMap() + for (superClassType in evaluateSuperClassesAsTypeHints(cls, context.typeContext)) { + RecursionManager.doPreventingRecursion( + superClassType.pyClass, + false) { + getGenericSubstitutions(superClassType.pyClass, context) + }?.let { superSubstitutions -> + results.putAll(superSubstitutions) } // TODO Share this logic with PyTypeChecker.collectTypeSubstitutions - List superTypeParameters = collectTypeParameters(superClassType.getPyClass(), context); - List superTypeArguments = superClassType instanceof PyCollectionType parameterized ? - parameterized.getElementTypes() : Collections.emptyList(); - PyTypeParameterMapping mapping = - PyTypeParameterMapping.mapByShape(superTypeParameters, superTypeArguments, Option.MAP_UNMATCHED_EXPECTED_TYPES_TO_ANY, - Option.USE_DEFAULTS); + val superTypeParameters = collectTypeParameters(superClassType.pyClass, context) + val superTypeArguments = if (superClassType is PyCollectionType) superClassType.elementTypes else mutableListOf() + val mapping = + PyTypeParameterMapping.mapByShape( + superTypeParameters, superTypeArguments, PyTypeParameterMapping.Option.MAP_UNMATCHED_EXPECTED_TYPES_TO_ANY, + PyTypeParameterMapping.Option.USE_DEFAULTS + ) if (mapping != null) { - for (Couple pair : mapping.getMappedTypes()) { - PyType expectedType = pair.getFirst(); - PyType actualType = pair.getSecond(); - if (!expectedType.equals(actualType)) { - results.put(expectedType, actualType); + for (pair in mapping.mappedTypes) { + val expectedType = pair.getFirst() + val actualType = pair.getSecond() + if (expectedType != actualType) { + results[expectedType] = actualType } } } } - return results; + return results } - private static @NotNull List evaluateSuperClassesAsTypeHints(@NotNull PyClass pyClass, @NotNull TypeEvalContext context) { - List results = new ArrayList<>(); - for (PyExpression superClassExpression : getSuperClassExpressions(pyClass)) { - PyType type = Ref.deref(getType(superClassExpression, context)); - if (type instanceof PyClassType classType) { - results.add(classType); - } - } - return results; - } - - private static @NotNull List collectTypeParameters(@NotNull PyClass cls, @NotNull Context context) { - if (!isGeneric(cls, context.getTypeContext())) { - return Collections.emptyList(); - } - if (cls.getTypeParameterList() != null) { - List typeParameters = cls.getTypeParameterList().getTypeParameters(); - return StreamEx.of(typeParameters) - .map(typeParameter -> getTypeParameterTypeFromTypeParameter(typeParameter, context)) - .nonNull() - .toList(); - } - // See https://mypy.readthedocs.io/en/stable/generics.html#defining-sub-classes-of-generic-classes - List parameterizedSuperClassExpressions = - ContainerUtil.filterIsInstance(getSuperClassExpressions(cls), PySubscriptionExpression.class); - PySubscriptionExpression genericAsSuperClass = ContainerUtil.find(parameterizedSuperClassExpressions, s -> { - return resolvesToQualifiedNames(s.getOperand(), context.myContext, GENERIC); - }); - return StreamEx.of(genericAsSuperClass != null ? Collections.singletonList(genericAsSuperClass) : parameterizedSuperClassExpressions) - .map(PySubscriptionExpression::getIndexExpression) - .flatMap(e -> { - final PyTupleExpression tupleExpr = as(e, PyTupleExpression.class); - return tupleExpr != null ? StreamEx.of(tupleExpr.getElements()) : StreamEx.of(e); - }) - .nonNull() - .map(e -> getType(e, context)) - .map(Ref::deref) - .flatMap(type -> { - PyTypeChecker.Generics typeParams = PyTypeChecker.collectGenerics(type, context.myContext); - return StreamEx.of(typeParams.getTypeVars()).append(typeParams.getTypeVarTuples()) - .append(StreamEx.of(typeParams.getParamSpecs())); - }) - .select(PyTypeParameterType.class) - .distinct() - .toList(); - } - - /** - * If the class' stub is present, return expressions in the base classes list, converting - * their saved text chunks into {@link PyExpressionCodeFragment} and extracting top-level expressions - * from them. Otherwise, get superclass expressions directly from AST. - */ - private static @NotNull List getSuperClassExpressions(@NotNull PyClass pyClass) { - final PyClassStub classStub = pyClass.getStub(); - if (classStub == null) { - return List.of(pyClass.getSuperClassExpressions()); - } - return CachedValuesManager.getCachedValue(pyClass, () -> CachedValueProvider.Result.create( - ContainerUtil.mapNotNull(classStub.getSuperClassesText(), x -> PyUtil.createExpressionFromFragment(x, pyClass)), - PsiModificationTracker.MODIFICATION_COUNT) - ); - } - - private static @NotNull List collectTypeParametersFromTypeAliasStatement(@NotNull PyTypeAliasStatement typeAliasStatement, - @NotNull Context context) { - PyTypeParameterList typeParameterList = typeAliasStatement.getTypeParameterList(); - if (typeParameterList != null) { - List typeParameters = typeParameterList.getTypeParameters(); - return StreamEx.of(typeParameters) - .map(typeParameter -> getTypeParameterTypeFromTypeParameter(typeParameter, context)) - .nonNull() - .toList(); - } - return Collections.emptyList(); - } - - public static boolean isGeneric(@NotNull PyWithAncestors descendant, @NotNull TypeEvalContext context) { - if (descendant instanceof PyClass pyClass && pyClass.getTypeParameterList() != null || - descendant instanceof PyClassType pyClassType && pyClassType.getPyClass().getTypeParameterList() != null) { - return true; - } - for (PyClassLikeType ancestor : descendant.getAncestorTypes(context)) { - if (ancestor != null) { - if (GENERIC_CLASSES.contains(ancestor.getClassQName())) { - return true; - } - else if (ancestor instanceof PyClassType classType && - classType.getPyClass().getTypeParameterList() != null) { - return true; - } - } - } - return false; - } - - @ApiStatus.Internal - public static @Nullable Ref getType(@NotNull PyExpression expression, @NotNull TypeEvalContext context, boolean useFqn) { - return staticWithCustomContext(context, useFqn, customContext -> getType(expression, customContext)); - } - - public static @Nullable Ref<@Nullable PyType> getType(@NotNull PyExpression expression, @NotNull TypeEvalContext context) { - return staticWithCustomContext(context, customContext -> getType(expression, customContext)); - } - - public static @Nullable Ref<@Nullable PyType> getTypeForTypeHint(@NotNull PyExpression expression, @NotNull TypeEvalContext context) { - return staticWithCustomContext(context, true, customContext -> getTypeForResolvedElement(expression, null, expression, customContext)); - } - - private static @Nullable Ref<@Nullable PyType> getType(@NotNull PyExpression expression, @NotNull Context context) { - PyType type = context.getKnownType(expression); - if (type != null) { - return Ref.create(type); - } - for (Pair pair : tryResolvingWithAliases(expression, context.getTypeContext())) { - final Ref typeRef = getTypeForResolvedElement(expression, pair.getFirst(), pair.getSecond(), context); - if (typeRef != null) { - if (typeRef.get() != null) { - context.assumeType(expression, typeRef.get()); - } - return typeRef; - } - } - return null; - } - - private static boolean typeHasOverloadedBitwiseOr(@NotNull PyType type, @NotNull PyExpression expression, - @NotNull Context context) { - if (!(type instanceof PyClassType classType)) { - return false; - } - TypeEvalContext typeContext = context.getTypeContext(); - PyClassLikeType metaClassType = classType.getMetaClassType(typeContext, true); - if (metaClassType == null) { - return false; - } - var resolved = metaClassType - .resolveMember("__or__", expression, AccessDirection.READ, PyResolveContext.defaultContext(typeContext)); - if (resolved == null || resolved.isEmpty()) return false; - - return StreamEx.of(resolved) - .map(it -> it.getElement()) - .nonNull() - .noneMatch(it -> PyBuiltinCache.getInstance(it).isBuiltin(it)); - } - - public static boolean isBitwiseOrUnionAvailable(@NotNull TypeEvalContext context) { - final PsiFile originFile = context.getOrigin(); - return originFile == null || isBitwiseOrUnionAvailable(originFile); - } - - public static boolean isBitwiseOrUnionAvailable(@NotNull PsiElement element) { - if (LanguageLevel.forElement(element).isAtLeast(LanguageLevel.PYTHON310)) return true; - - PsiFile file = element.getContainingFile(); - if (file instanceof PyFile && ((PyFile)file).hasImportFromFuture(FutureFeature.ANNOTATIONS)) { - return file == element || PsiTreeUtil.getParentOfType(element, PyAnnotation.class, false, PyStatement.class) != null; - } - - return false; - } - - private static @Nullable Ref getTypeForResolvedElement(@NotNull PyExpression typeHint, - @Nullable PyQualifiedNameOwner alias, - @NotNull PsiElement resolved, - @NotNull Context context) { - if (alias != null) { - if (context.containsTypeAlias(alias)) { - // Recursive types are not yet supported - return null; - } - context.addTypeAlias(alias); - } - if (resolved instanceof PyClass pyClass && !context.addClassDeclaration(pyClass)) { - // Resolving to normal classes shouldn't cause recursive evaluation of type hints, - // but constructing recursive PyTypedDictTypes will trigger that. - return null; - } - try { - final Ref typeHintFromProvider = PyTypeHintProvider.Companion.parseTypeHint( - typeHint, - alias, - resolved, - context.getTypeContext() - ); - if (typeHintFromProvider != null) { - return typeHintFromProvider; - } - final Ref typeFromParenthesizedExpression = getTypeFromParenthesizedExpression(resolved, context); - if (typeFromParenthesizedExpression != null) { - return typeFromParenthesizedExpression; - } - // We perform chained resolve only for actual aliases as tryResolvingWithAliases() returns the passed-in - // expression both when it's not a reference expression and when it's failed to resolve it, hence we might - // hit SOE for mere unresolved references in the latter case. - if (alias != null) { - Ref typeFromTypeAlias = getTypeFromTypeAlias(alias, typeHint, resolved, context); - if (typeFromTypeAlias != null) { - return typeFromTypeAlias; - } - } - final PyType neverType = getNeverType(resolved); - if (neverType != null) { - return Ref.create(neverType); - } - final Ref unionType = getUnionType(resolved, context); - if (unionType != null) { - return unionType; - } - final Ref intersectionType = getIntersectionType(resolved, context); - if (intersectionType != null) { - return intersectionType; - } - final PyType concatenateType = getConcatenateType(resolved, context); - if (concatenateType != null) { - return Ref.create(concatenateType); - } - final Ref optionalType = getOptionalType(resolved, context); - if (optionalType != null) { - return optionalType; - } - final PyType callableType = getCallableType(resolved, context); - if (callableType != null) { - return Ref.create(callableType); - } - final Ref classVarType = unwrapTypeModifier(resolved, context, CLASS_VAR); - if (classVarType != null) { - return classVarType; - } - final Ref classObjType = getClassObjectType(resolved, context); - if (classObjType != null) { - return classObjType; - } - final Ref finalType = unwrapTypeModifier(resolved, context, FINAL, FINAL_EXT); - if (finalType != null) { - return finalType; - } - final Ref annotatedType = getAnnotatedType(resolved, context); - if (annotatedType != null) { - return annotatedType; - } - final Ref requiredOrNotRequiredType = getTypedDictSpecialItemType(resolved, context); - if (requiredOrNotRequiredType != null) { - return requiredOrNotRequiredType; - } - final Ref literalStringType = getLiteralStringType(resolved, context); - if (literalStringType != null) { - return literalStringType; - } - final Ref literalType = getLiteralType(resolved, context); - if (literalType != null) { - return literalType; - } - final Ref typeAliasType = getExplicitTypeAliasType(resolved); - if (typeAliasType != null) { - return typeAliasType; - } - final Ref narrowedType = getNarrowedType(resolved, context); - if (narrowedType != null) { - return narrowedType; - } - final PyType parameterizedType = getParameterizedType(resolved, context); - if (parameterizedType != null) { - return Ref.create(parameterizedType); - } - final PyType collection = getCollection(resolved, context.getTypeContext()); - if (collection != null) { - return Ref.create(collection); - } - final PyType typeParameter = getTypeParameterTypeFromDeclaration(resolved, context); - if (typeParameter != null) { - return Ref.create(anchorTypeParameter(typeHint, typeParameter, context)); - } - final PyType unpackedType = getUnpackedType(resolved, context.getTypeContext()); - if (unpackedType != null) { - return Ref.create(unpackedType); - } - final PyType typeParameterType = getTypeParameterTypeFromTypeParameter(resolved, context); - if (typeParameterType != null) { - return Ref.create(typeParameterType); - } - final PyType callableParameterListType = getCallableParameterListType(resolved, context); - if (callableParameterListType != null) { - return Ref.create(callableParameterListType); - } - final PyType stringBasedType = getStringLiteralType(resolved, context); - if (stringBasedType != null) { - return Ref.create(stringBasedType); - } - final Ref anyType = getAnyType(resolved, context); - if (anyType != null) { - return anyType; - } - final PyType typedDictType = PyTypedDictTypeProvider.Companion.getTypedDictTypeForResolvedElement(resolved, context.getTypeContext()); - if (typedDictType != null) { - return Ref.create(typedDictType); - } - final Ref selfType = getSelfType(resolved, typeHint, context); - if (selfType != null) { - return selfType; - } - final Ref noneType = getNoneType(typeHint, resolved); - if (noneType != null) { - return noneType; - } - PyTypingNewType newType = PyTypingNewTypeTypeProvider.getNewTypeForResolvedElement(resolved, context.getTypeContext()); - if (newType != null) { - return Ref.create(newType.toInstance()); - } - final Ref classType = getClassType(typeHint, resolved, context); - if (classType != null) { - return classType; - } - if (context.myUseFqn) { - if (resolved.getText().equals("Unknown")) { - return Ref.create(); - } - if (resolved instanceof PyFunctionTypeRepresentation function) { - var result = context.myContext.getType(function); - if (result != null) { - return Ref.create(result); - } - } - if (resolved instanceof PySubscriptionExpression subscriptionExpression) { - var moduleType = getModuleType(subscriptionExpression); - if (moduleType != null) { - return moduleType; - } - } - } - return null; - } - finally { - if (resolved instanceof PyClass pyClass) { - context.removeClassDeclaration(pyClass); - } - if (alias != null) { - context.removeTypeAlias(alias); - } - } - } - - private static @Nullable Ref getModuleType(PySubscriptionExpression moduleDefinition) { - String name = moduleDefinition.getRootOperand().getName(); - if (name == null || !name.equals(PyModuleTypeName)) { - return null; - } - if (!(moduleDefinition.getIndexExpression() instanceof PyReferenceExpression moduleReferenceExpression)) { - return null; - } - var moduleName = moduleReferenceExpression.getName(); - if (moduleName == null) { - return null; - } - Project project = moduleDefinition.getProject(); - var moduleInitFiles = PyModuleNameIndex.findByShortName(moduleName, project, GlobalSearchScope.everythingScope(project)); - var skeletons = PythonSdkUtil.getSkeletonsRootPath(PathManager.getSystemDir().toString()); - var firstModuleInitFile = ContainerUtil.find(moduleInitFiles, a -> a != null && !Path.of(a.getVirtualFile().getPath()).startsWith(skeletons)); - if (firstModuleInitFile == null) { - firstModuleInitFile = ContainerUtil.find(moduleInitFiles, a -> a != null); - } - - if (firstModuleInitFile == null) { - return null; - } - return Ref.create(new PyModuleType(firstModuleInitFile)); - } - - private static Ref getIntersectionType(@NotNull PsiElement resolved, @NotNull PyTypingTypeProvider.Context context) { - if (resolved instanceof PyBinaryExpression expression && expression.getOperator() == PyTokenTypes.AND) { - PyExpression left = expression.getLeftExpression(); - PyExpression right = expression.getRightExpression(); - if (left == null || right == null) return null; - - Ref leftTypeRef = getType(left, context); - Ref rightTypeRef = getType(right, context); - if (leftTypeRef == null || rightTypeRef == null) return null; - - PyType intersection = PyIntersectionType.intersection(leftTypeRef.get(), rightTypeRef.get()); - return intersection != null ? Ref.create(intersection) : null; - } - return null; - } - - private static @Nullable Ref getNoneType(@NotNull PyExpression typeHint, @NotNull PsiElement resolved) { - if (typeHint instanceof PyNoneLiteralExpression || - typeHint instanceof PyReferenceExpression && PyNames.NONE.equals(typeHint.getText())) { - return Ref.create(PyBuiltinCache.getInstance(resolved).getNoneType()); - } - return null; - } - - private static @Nullable PyType getCallableParameterListType(@NotNull PsiElement resolved, @NotNull Context context) { - if (resolved instanceof PyListLiteralExpression listLiteral) { - List argumentTypes = ContainerUtil.map(listLiteral.getElements(), defExpr -> Ref.deref(getType(defExpr, context))); - return new PyCallableParameterListTypeImpl(ContainerUtil.map(argumentTypes, PyCallableParameterImpl::nonPsi)); - } - return null; - } - - private static @Nullable Ref getNarrowedType(@NotNull PsiElement resolved, @NotNull Context context) { - if (resolved instanceof PySubscriptionExpression subscriptionExpr) { - Collection names = resolveToQualifiedNames(subscriptionExpr.getOperand(), context.getTypeContext()); - var isTypeIs = names.contains(TYPE_IS) || names.contains(TYPE_IS_EXT); - var isTypeGuard = names.contains(TYPE_GUARD) || names.contains(TYPE_GUARD_EXT); - if (isTypeIs || isTypeGuard) { - List indexTypes = getIndexTypes(subscriptionExpr, context); - if (indexTypes.size() == 1) { - PyNarrowedType narrowedType = PyNarrowedType.Companion.create(subscriptionExpr, isTypeIs, indexTypes.get(0)); - if (narrowedType != null) { - return Ref.create(narrowedType); - } - } - } - } - return null; - } - - private static Ref getSelfType(@NotNull PsiElement resolved, @NotNull PyExpression typeHint, @NotNull Context context) { - if (resolved instanceof PyQualifiedNameOwner && - (SELF.equals(((PyQualifiedNameOwner)resolved).getQualifiedName()) || - SELF_EXT.equals(((PyQualifiedNameOwner)resolved).getQualifiedName()))) { - PsiElement typeHintContext = getStubRetainedTypeHintContext(typeHint); - - PyClass containingClass = typeHintContext instanceof PyClass ? (PyClass)typeHintContext - : PsiTreeUtil.getStubOrPsiParentOfType(typeHintContext, PyClass.class); - if (containingClass == null) return null; - - PyClassType scopeClassType = as(containingClass.getType(context.getTypeContext()), PyClassType.class); - if (scopeClassType == null) return null; - - return Ref.create(new PySelfType(scopeClassType).toInstance()); - } - return null; - } - - private static @Nullable Ref getTypeFromParenthesizedExpression(@NotNull PsiElement resolved, @NotNull Context context) { - if (resolved instanceof PyParenthesizedExpression) { - final PyExpression containedExpression = PyPsiUtils.flattenParens((PyExpression)resolved); - return containedExpression != null ? getType(containedExpression, context) : null; - } - return null; - } - - private static @Nullable Ref getExplicitTypeAliasType(@NotNull PsiElement resolved) { - if (resolved instanceof PyQualifiedNameOwner) { - String qualifiedName = ((PyQualifiedNameOwner)resolved).getQualifiedName(); - if (TYPE_ALIAS.equals(qualifiedName) || TYPE_ALIAS_EXT.equals(qualifiedName)) { - return Ref.create(); - } - } - return null; - } - - private static @Nullable PyType anchorTypeParameter(@NotNull PyExpression typeHint, @Nullable PyType type, @NotNull Context context) { - PyQualifiedNameOwner typeParamDefinitionFromStack = context.isTypeAliasStackEmpty() ? null : context.peekTypeAlias(); - assert typeParamDefinitionFromStack == null || typeParamDefinitionFromStack instanceof PyTargetExpression; - PyTargetExpression targetExpr = (PyTargetExpression)typeParamDefinitionFromStack; - if (type instanceof PyTypeVarTypeImpl typeVar) { - return typeVar.withScopeOwner(getTypeParameterScope(typeVar.getName(), typeHint, context)).withDeclarationElement(targetExpr); - } - if (type instanceof PyParamSpecType paramSpec) { - return paramSpec.withScopeOwner(getTypeParameterScope(paramSpec.getName(), typeHint, context)).withDeclarationElement(targetExpr); - } - if (type instanceof PyTypeVarTupleTypeImpl typeVarTuple) { - return typeVarTuple.withScopeOwner(getTypeParameterScope(typeVarTuple.getName(), typeHint, context)) - .withDeclarationElement(targetExpr); - } - return type; - } - - private static @Nullable Ref getClassObjectType(@NotNull PsiElement resolved, @NotNull Context context) { - if (resolved instanceof PySubscriptionExpression subsExpr) { - final PyExpression operand = subsExpr.getOperand(); - if (resolvesToQualifiedNames(operand, context.getTypeContext(), TYPE, PyNames.TYPE)) { - final PyExpression indexExpr = subsExpr.getIndexExpression(); - if (indexExpr != null) { - if (resolvesToQualifiedNames(indexExpr, context.getTypeContext(), ANY)) { - return Ref.create(PyBuiltinCache.getInstance(resolved).getTypeType()); - } - return getAsClassObjectType(indexExpr, context); - } - // Map Type[Something] with unsupported type parameter to Any, instead of generic type for the class "type" - return Ref.create(); - } - } - // Replace plain non-parametrized Type with its builtin counterpart - else if (TYPE.equals(getQualifiedName(resolved))) { - return Ref.create(PyBuiltinCache.getInstance(resolved).getTypeType()); - } - return null; - } - - private static @NotNull Ref getAsClassObjectType(@NotNull PyExpression expression, @NotNull Context context) { - final PyType type = Ref.deref(getType(expression, context)); - final PyClassType classType = as(type, PyClassType.class); - if (classType != null && !classType.isDefinition()) { - return Ref.create(classType.toClass()); - } - final PyTypeVarType typeVar = as(type, PyTypeVarType.class); - if (typeVar != null && !typeVar.isDefinition()) { - return Ref.create(typeVar.toClass()); - } - final PySelfType selfType = as(type, PySelfType.class); - if (selfType != null) { - return Ref.create(selfType.toClass()); - } - // Represent Type[Union[str, int]] internally as Union[Type[str], Type[int]] - if (type instanceof PyUnionType unionType && - ContainerUtil.all(unionType.getMembers(), t -> t instanceof PyClassType clsType && !clsType.isDefinition())) { - //noinspection DataFlowIssue - return Ref.create(unionType.map(clsType -> ((PyClassType)clsType).toClass())); - } - return Ref.create(); - } - - private static @Nullable Ref getAnyType(@NotNull PsiElement element, @NotNull Context context) { - if (ANY.equals(getQualifiedName(element))) { - return Ref.create(); - } - if (context.myUseFqn && ANY.equals(element.getText())) { - return Ref.create(); - } - return null; - } - - private static @Nullable Ref getClassType(@NotNull PyExpression typeHint, @NotNull PsiElement element, @NotNull Context context) { - if (typeHint instanceof PyReferenceExpression && element instanceof PyTypedElement) { - TypeEvalContext typeContext = context.getTypeContext(); - final PyType type; - if (context.myUseFqn && element instanceof PyReferenceExpression referenceExpression) { - var qualifiedName = referenceExpression.asQualifiedName(); - var project = element.getProject(); - var class_ = qualifiedName != null ? PyPsiFacade.getInstance(project).createClassByQName(qualifiedName.toString(), element) : null; - type = class_ != null ? class_.getType(typeContext) : null; - } - else { - type = typeContext.getType((PyTypedElement)element); - } - if (type instanceof PyClassLikeType classLikeType) { - if (classLikeType.isDefinition()) { - // If we're interpreting a type hint like "MyGeneric" that is not followed by a list of type arguments (e.g. MyGeneric[int]), - // we want to parameterize it with its type parameters defaults already here. - // We need this check for the type argument list because getParameterizedType() relies on getClassType() for - // getting the type corresponding to the subscription expression operand. - PsiElement stubRetainedContext = getStubRetainedTypeHintContext(typeHint); - if (classLikeType instanceof PyClassType classType && - !(stubRetainedContext instanceof PyClass || - PsiTreeUtil.getStubOrPsiParentOfType(stubRetainedContext, ScopeOwner.class) instanceof PyClass) && - !(typeHint.getParent() instanceof PySubscriptionExpression se && typeHint.equals(se.getOperand())) && - isGeneric(classType, context.myContext)) { - PyCollectionType parameterized = parameterizeClassDefaultAware(classType.getPyClass(), List.of(), context); - if (parameterized != null) { - return Ref.create(parameterized.toInstance()); - } - } - final PyType instanceType = classLikeType.toInstance(); - return Ref.create(instanceType); - } - } - } - return null; - } - - private static @Nullable Ref getOptionalType(@NotNull PsiElement element, @NotNull Context context) { - if (element instanceof PySubscriptionExpression subscriptionExpr) { - if (resolvesToQualifiedNames(subscriptionExpr.getOperand(), context.getTypeContext(), OPTIONAL)) { - final PyExpression indexExpr = subscriptionExpr.getIndexExpression(); - if (indexExpr != null) { - final Ref typeRef = getType(indexExpr, context); - if (typeRef != null) { - return Ref.create(PyUnionType.union(typeRef.get(), PyBuiltinCache.getInstance(element).getNoneType())); - } - } - return Ref.create(); - } - } - return null; - } - - private static @Nullable Ref getLiteralStringType(@NotNull PsiElement resolved, @NotNull Context context) { - if (resolved instanceof PyTargetExpression referenceExpression) { - if (resolvesToQualifiedNames(referenceExpression, context.getTypeContext(), LITERALSTRING, LITERALSTRING_EXT)) { - return Ref.create(PyLiteralStringType.Companion.create(resolved)); - } - } - - return null; - } - - private static @Nullable Ref getLiteralType(@NotNull PsiElement resolved, @NotNull Context context) { - if (resolved instanceof PySubscriptionExpression subscriptionExpr) { - if (resolvesToQualifiedNames(subscriptionExpr.getOperand(), context, LITERAL, LITERAL_EXT)) { - return Optional - .ofNullable(subscriptionExpr.getIndexExpression()) - .map(index -> PyLiteralType.Companion.fromLiteralParameter(index, context.getTypeContext())) - .map(Ref::create) - .orElse(null); - } - } - - return null; - } - - private static @Nullable Ref getAnnotatedType(@NotNull PsiElement resolved, @NotNull Context context) { - if (resolved instanceof PySubscriptionExpression subscriptionExpr) { - final PyExpression operand = subscriptionExpr.getOperand(); - if (resolvesToQualifiedNames(operand, context.getTypeContext(), ANNOTATED, ANNOTATED_EXT)) { - final PyExpression indexExpr = subscriptionExpr.getIndexExpression(); - final PyExpression type = indexExpr instanceof PyTupleExpression ? ((PyTupleExpression)indexExpr).getElements()[0] : indexExpr; - if (type != null) { - return getType(type, context); - } - } - } - - return null; - } - - private static @Nullable Ref getTypedDictSpecialItemType(@NotNull PsiElement resolved, @NotNull Context context) { - if (resolved instanceof PySubscriptionExpression subscriptionExpr) { - final PyExpression operand = subscriptionExpr.getOperand(); - if (resolvesToQualifiedNames(operand, context.getTypeContext(), - REQUIRED, REQUIRED_EXT, - NOT_REQUIRED, NOT_REQUIRED_EXT, - READONLY, READONLY_EXT)) { - final PyExpression indexExpr = subscriptionExpr.getIndexExpression(); - final PyExpression type = indexExpr instanceof PyTupleExpression ? ((PyTupleExpression)indexExpr).getElements()[0] : indexExpr; - if (type != null) { - return getType(type, context); - } - } - } - - return null; - } - - private static @Nullable Ref unwrapTypeModifier(@NotNull PsiElement resolved, @NotNull Context context, String... type) { - if (resolved instanceof PySubscriptionExpression subscriptionExpr) { - if (resolvesToQualifiedNames(subscriptionExpr.getOperand(), context.getTypeContext(), type)) { - final PyExpression indexExpr = subscriptionExpr.getIndexExpression(); - if (indexExpr != null) { - return getType(indexExpr, context); - } - } - } - - return null; - } - - private static boolean typeHintedWithName(@NotNull T owner, - @NotNull TypeEvalContext context, - String... names) { - return ContainerUtil.exists(names, resolveTypeHintsToQualifiedNames(owner, context)::contains); - } - - private static Collection resolveTypeHintsToQualifiedNames( - @NotNull T owner, - @NotNull TypeEvalContext context + @JvmRecord + data class GeneratorTypeDescriptor( + @JvmField val yieldType: PyType?, + @JvmField val sendType: PyType?, + @JvmField val returnType: PyType?, + @JvmField val isAsync: Boolean, ) { - var annotation = getAnnotationValue(owner, context); - if (annotation instanceof PyStringLiteralExpression stringLiteralExpression) { - final var annotationText = stringLiteralExpression.getStringValue(); - annotation = PyUtil.createExpressionFromFragment(annotationText, owner); - if (annotation == null) return Collections.emptyList(); - } + companion object { + /** + * Extracts type parameters from typing.Generator and typing.AsyncGenerator + */ + fun fromGenerator(type: PyType?): GeneratorTypeDescriptor? { + if (type !is PyClassType) return null - if (annotation instanceof PySubscriptionExpression pySubscriptionExpression) { - return resolveToQualifiedNames(pySubscriptionExpression.getOperand(), context); - } - else if (annotation instanceof PyReferenceExpression) { - return resolveToQualifiedNames(annotation, context); - } + val qName = type.classQName + if (qName == null) return null - final String typeCommentValue = owner.getTypeCommentAnnotation(); - final PyExpression typeComment = typeCommentValue == null ? null : PyUtil.createExpressionFromFragment(typeCommentValue, owner); - if (typeComment instanceof PySubscriptionExpression pySubscriptionExpression) { - return resolveToQualifiedNames(pySubscriptionExpression.getOperand(), context); - } - else if (typeComment instanceof PyReferenceExpression) { - return resolveToQualifiedNames(typeComment, context); - } + val isAsync = ASYNC_GENERATOR == qName + if (!isAsync && GENERATOR != qName) return null - return Collections.emptyList(); + val noneType: PyType? = getInstance(type.pyClass).noneType + + var yieldType: PyType? = null + var sendType = noneType + var returnType = if (isAsync) null else noneType + if (type is PyCollectionType) { + yieldType = type.elementTypes.getOrNull(0) + sendType = type.elementTypes.getOrElse(1) { sendType } + returnType = type.elementTypes.getOrElse(2) { returnType } + } + return GeneratorTypeDescriptor(yieldType, sendType, returnType, isAsync) + } + + /** + * Unlike [.fromGenerator], this method can also extract yield type from Protocol types like typing.Iterable + */ + @JvmStatic + fun fromGeneratorOrProtocol(type: PyType?, context: TypeEvalContext): GeneratorTypeDescriptor? { + if (type !is PyClassType) return null + + val desc = fromGenerator(type) + if (desc != null) { + return desc + } + + if (type.isProtocol(context)) { + var yieldType: PyType? + + val syncUpcast = PyTypeUtil.convertToType(type, "typing.Iterable", type.pyClass, context) + if (syncUpcast is PyCollectionType) { + yieldType = syncUpcast.iteratedItemType + return GeneratorTypeDescriptor(yieldType, null, null, false) + } + val asyncUpcast = PyTypeUtil.convertToType(type, "typing.AsyncIterable", type.pyClass, context) + if (asyncUpcast is PyCollectionType) { + yieldType = asyncUpcast.iteratedItemType + return GeneratorTypeDescriptor(yieldType, null, null, true) + } + + // Here we try to understand a yield type by return type of __next__ method of protocol specified in annotation. + // We cannot use convertToType with typing.Iterator here, as it inherits from typing.Iterable + // and requires both __iter__ and __next__, while it should be possible to decide the yield type only by __next__. + // TODO: unify logic with PyTargetExpressionImpl.getIterationType (PY-82453) + val next = type.pyClass.findMethodByName(PyNames.DUNDER_NEXT, true, context) + if (next != null) { + yieldType = context.getReturnType(next) + yieldType = PyTypeChecker.substitute(yieldType, PyTypeChecker.unifyReceiver(type, context), context) + return GeneratorTypeDescriptor(yieldType, null, null, false) + } + + val anext = type.pyClass.findMethodByName(PyNames.ANEXT, true, context) + if (anext != null) { + yieldType = Ref.deref(unwrapCoroutineReturnType(context.getReturnType(anext))) + yieldType = PyTypeChecker.substitute(yieldType, PyTypeChecker.unifyReceiver(type, context), context) + return GeneratorTypeDescriptor(yieldType, null, null, true) + } + } + return null + } + } } - public static boolean isFinal(@NotNull PyDecoratable decoratable, @NotNull TypeEvalContext context) { - return ContainerUtil.exists(PyKnownDecoratorUtil.getKnownDecorators(decoratable, context), - d -> d == TYPING_FINAL || d == TYPING_FINAL_EXT); - } - - public static boolean isFinal(@NotNull T owner, @NotNull TypeEvalContext context) { - return PyUtil.getParameterizedCachedValue(owner, context, p -> - typeHintedWithName(owner, context, FINAL, FINAL_EXT)); - } - - public static boolean isClassVar(@NotNull T owner, @NotNull TypeEvalContext context) { - return PyUtil.getParameterizedCachedValue(owner, context, p -> - typeHintedWithName(owner, context, CLASS_VAR)); - } - - private static boolean resolvesToQualifiedNames(@NotNull PyExpression expression, @NotNull Context context, String... names) { - if (!context.myUseFqn) return resolvesToQualifiedNames(expression, context.myContext, names); - if (!(expression instanceof PyReferenceExpression referenceExpression)) return false; - var qualifier = referenceExpression.getQualifier(); - if (qualifier == null) return false; - var qName = qualifier.getName() + "." + expression.getName(); - return ContainerUtil.exists(names, name -> name.equals(qName)); - } - - private static boolean resolvesToQualifiedNames(@NotNull PyExpression expression, @NotNull TypeEvalContext context, String... names) { - final var qualifiedNames = resolveToQualifiedNames(expression, context); - return ContainerUtil.exists(names, qualifiedNames::contains); + override fun withCustomContext(context: TypeEvalContext, delegate: java.util.function.Function): T? { + return staticWithCustomContext(context, delegate::apply) } @ApiStatus.Internal - public static @Nullable PyExpression getAnnotationValue(@NotNull PyAnnotationOwner owner, @NotNull TypeEvalContext context) { - if (context.maySwitchToAST(owner)) { - final PyAnnotation annotation = owner.getAnnotation(); - if (annotation != null) { - return annotation.getValue(); - } - } - else { - final String annotationText = owner.getAnnotationValue(); - if (annotationText != null) { - return PyUtil.createExpressionFromFragment(annotationText, owner); - } - } - return null; - } - - public static @Nullable Ref getStringBasedType(@NotNull String contents, - @NotNull PsiElement anchor, - @NotNull TypeEvalContext context) { - return staticWithCustomContext(context, c -> getStringBasedType(contents, anchor, c)); - } - - private static @Nullable Ref getStringBasedType(@NotNull String contents, @NotNull PsiElement anchor, @NotNull Context context) { - return doPreventingRecursion(Pair.create(anchor, contents), true, () -> { - final PyExpression expr = PyUtil.createExpressionFromFragment(contents, anchor); - return expr != null ? getType(expr, context) : null; - }); - } - - private static @Nullable PyType getStringLiteralType(@NotNull PsiElement element, @NotNull Context context) { - if (element instanceof PyStringLiteralExpression stringLiteral) { - final String contents = stringLiteral.getStringValue(); - // A multiline string literal can contain a type expression unparsable without parentheses - return Ref.deref(getStringBasedType(contents.contains("\n") ? "(" + contents + ")" : contents, element, context)); - } - return null; - } - - private static @Nullable Ref getVariableTypeCommentType(@NotNull String contents, - @NotNull PsiElement element, - @NotNull Context context) { - final PyExpression expr = PyPsiUtils.flattenParens(PyUtil.createExpressionFromFragment(contents, element)); - if (expr != null) { - if (element instanceof PyTargetExpression target && expr instanceof PyTupleExpression) { - // Such syntax is specific to "# type:" comments, unpacking in type hints is not allowed anywhere else - // XXX: Switches stub to AST - final PyExpression topmostTarget = findTopmostTarget(target); - if (topmostTarget != null) { - final Map targetToExpr = mapTargetsToAnnotations(topmostTarget, expr); - final PyExpression typeExpr = targetToExpr.get(target); - if (typeExpr != null) { - return getType(typeExpr, context); - } - } - } - else { - return getType(expr, context); - } - } - return null; - } - - private static @Nullable PyExpression findTopmostTarget(@NotNull PyTargetExpression target) { - final PyElement validTargetParent = PsiTreeUtil.getParentOfType(target, PyForPart.class, PyWithItem.class, PyAssignmentStatement.class); - if (validTargetParent == null) { - return null; - } - final PyExpression topmostTarget = as(PsiTreeUtil.findPrevParent(validTargetParent, target), PyExpression.class); - if (validTargetParent instanceof PyForPart && topmostTarget != ((PyForPart)validTargetParent).getTarget()) { - return null; - } - if (validTargetParent instanceof PyWithItem && topmostTarget != ((PyWithItem)validTargetParent).getTarget()) { - return null; - } - if (validTargetParent instanceof PyAssignmentStatement && - ArrayUtil.indexOf(((PyAssignmentStatement)validTargetParent).getRawTargets(), topmostTarget) < 0) { - return null; - } - return topmostTarget; - } - - public static @NotNull Map mapTargetsToAnnotations(@NotNull PyExpression targetExpr, - @NotNull PyExpression typeExpr) { - final PyExpression targetsNoParen = PyPsiUtils.flattenParens(targetExpr); - final PyExpression typesNoParen = PyPsiUtils.flattenParens(typeExpr); - if (targetsNoParen == null || typesNoParen == null) { - return Collections.emptyMap(); - } - if (targetsNoParen instanceof PySequenceExpression && typesNoParen instanceof PySequenceExpression) { - final Ref> result = new Ref<>(new LinkedHashMap<>()); - mapTargetsToExpressions((PySequenceExpression)targetsNoParen, (PySequenceExpression)typesNoParen, result); - return result.isNull() ? Collections.emptyMap() : Collections.unmodifiableMap(result.get()); - } - else if (targetsNoParen instanceof PyTargetExpression && !(typesNoParen instanceof PySequenceExpression)) { - return ImmutableMap.of((PyTargetExpression)targetsNoParen, typesNoParen); - } - return Collections.emptyMap(); - } - - private static void mapTargetsToExpressions(@NotNull PySequenceExpression targetSequence, - @NotNull PySequenceExpression valueSequence, - @NotNull Ref> result) { - final PyExpression[] targets = targetSequence.getElements(); - final PyExpression[] values = valueSequence.getElements(); - - if (targets.length != values.length) { - result.set(null); - return; - } - - for (int i = 0; i < targets.length; i++) { - final PyExpression target = PyPsiUtils.flattenParens(targets[i]); - final PyExpression value = PyPsiUtils.flattenParens(values[i]); - - if (target == null || value == null) { - result.set(null); - return; - } - - if (target instanceof PySequenceExpression && value instanceof PySequenceExpression) { - mapTargetsToExpressions((PySequenceExpression)target, (PySequenceExpression)value, result); - if (result.isNull()) { - return; - } - } - else if (target instanceof PyTargetExpression && !(value instanceof PySequenceExpression)) { - final Map map = result.get(); - assert map != null; - map.put((PyTargetExpression)target, value); - } - else { - result.set(null); - return; - } - } - } - - private static @Nullable PyType getCallableType(@NotNull PsiElement resolved, @NotNull Context context) { - if (resolved instanceof PySubscriptionExpression subscriptionExpr) { - if (resolvesToQualifiedNames(subscriptionExpr.getOperand(), context.getTypeContext(), CALLABLE, CALLABLE_EXT)) { - final PyExpression indexExpr = subscriptionExpr.getIndexExpression(); - if (indexExpr instanceof PyTupleExpression tupleExpr) { - final PyExpression[] elements = tupleExpr.getElements(); - if (elements.length == 2) { - final PyExpression parametersExpr = elements[0]; - final PyExpression returnTypeExpr = elements[1]; - PyType returnType = Ref.deref(getType(returnTypeExpr, context)); - if (returnType instanceof PyVariadicType) { - returnType = null; - } - if (parametersExpr instanceof PyEllipsisLiteralExpression) { - return new PyCallableTypeImpl((PyCallableParameterVariadicType)null, returnType); - } - PyType parametersType = Ref.deref(getType(parametersExpr, context)); - if (parametersType instanceof PyCallableParameterVariadicType variadicType) { - return new PyCallableTypeImpl(variadicType, returnType); - } - } - } - } - } - else if (resolved instanceof PyTargetExpression targetExpression) { - if (resolvesToQualifiedNames(targetExpression, context.getTypeContext(), CALLABLE, CALLABLE_EXT)) { - return new PyCallableTypeImpl((PyCallableParameterVariadicType)null, null); - } - } - return null; - } - - private static @Nullable PyType getNeverType(@NotNull PsiElement element) { - var qName = getQualifiedName(element); - if (qName == null) return null; - if (List.of(NEVER, NEVER_EXT).contains(qName)) { - return PyNeverType.NEVER; - } - if (List.of(NO_RETURN, NO_RETURN_EXT).contains(qName)) { - return PyNeverType.NO_RETURN; - } - return null; - } - - private static @Nullable Ref getUnionType(@NotNull PsiElement element, @NotNull Context context) { - if (element instanceof PySubscriptionExpression subscriptionExpr) { - if (resolvesToQualifiedNames(subscriptionExpr.getOperand(), context.getTypeContext(), UNION)) { - PyType union = PyUnionType.union(getIndexTypes(subscriptionExpr, context)); - return union != null ? Ref.create(union) : null; - } - } - else if (element instanceof PyBinaryExpression expression && expression.getOperator() == PyTokenTypes.OR) { - PyExpression left = expression.getLeftExpression(); - PyExpression right = expression.getRightExpression(); - if (left == null || right == null) return null; - - Ref leftTypeRef = getType(left, context); - Ref rightTypeRef = getType(right, context); - if (leftTypeRef == null || rightTypeRef == null) return null; - - PyType leftType = leftTypeRef.get(); - if (leftType != null && typeHasOverloadedBitwiseOr(leftType, left, context)) return null; - - PyType union = PyUnionType.union(leftType, rightTypeRef.get()); - return union != null ? Ref.create(union) : null; - } - return null; - } - - private static @Nullable PyType getConcatenateType(@NotNull PsiElement element, @NotNull Context context) { - if (!(element instanceof PySubscriptionExpression subscriptionExpr)) return null; - if (!resolvesToQualifiedNames(subscriptionExpr.getOperand(), context.myContext, CONCATENATE, CONCATENATE_EXT)) return null; - if (!(subscriptionExpr.getIndexExpression() instanceof PyTupleExpression tupleExpression)) return null; - - List arguments = Arrays.asList(tupleExpression.getElements()); - if (arguments.size() < 2) return null; - - List prefixTypeExprs = arguments.subList(0, arguments.size() - 1); - List prefixTypes = ContainerUtil.map(prefixTypeExprs, it -> Ref.deref(getType(it, context.myContext))); - - PyExpression lastTypeExpr = arguments.get(arguments.size() - 1); - PyParamSpecType paramSpecType; - if (lastTypeExpr instanceof PyEllipsisLiteralExpression) { - paramSpecType = null; - } - else if (Ref.deref(getType(lastTypeExpr, context.myContext)) instanceof PyParamSpecType ps) { - paramSpecType = ps; - } - else { - return null; - } - return new PyConcatenateType(prefixTypes, paramSpecType); - } - - private static @Nullable PyTypeParameterType getTypeParameterTypeFromDeclaration(@NotNull PsiElement element, @NotNull Context context) { - if (element instanceof PyCallExpression assignedCall) { - PyAstTypeParameter.Kind typeParameterKind = getTypeParameterKindFromDeclaration(assignedCall, context.getTypeContext()); - if (typeParameterKind != null) { - final PyExpression[] arguments = assignedCall.getArguments(); - if (arguments.length > 0 && arguments[0] instanceof PyStringLiteralExpression nameArgument) { - final String name = nameArgument.getStringValue(); - PyExpression defaultExpression = assignedCall.getKeywordArgument("default"); - Ref defaultType = defaultExpression != null ? getType(defaultExpression, context) : null; - switch (typeParameterKind) { - case TypeVarTuple -> { - return new PyTypeVarTupleTypeImpl(name) - .withDefaultType(defaultType != null && Ref.deref(defaultType) instanceof PyPositionalVariadicType posVariadic - ? Ref.create(posVariadic) : null); - } - case TypeVar -> { - // TypeVar __init__ parameters: - // (name, *constraints, bound = None, contravariant = False, covariant = False, infer_variance = False, default = ...) - List constraints = Stream.of(arguments) - .skip(1) - .takeWhile(expr -> !(expr instanceof PyKeywordArgument)) - .map(expr -> Ref.deref(getType(expr, context))) - .toList(); - PyExpression boundExpression = assignedCall.getKeywordArgument("bound"); - PyType bound = boundExpression == null ? null : Ref.deref(getType(boundExpression, context)); - PyTypeVarType.Variance variance = getTypeVarVarianceFromDeclaration(assignedCall); - return new PyTypeVarTypeImpl(name, constraints, bound, defaultType, variance); - } - case ParamSpec -> { - return new PyParamSpecType(name) - .withDefaultType(defaultType != null && Ref.deref(defaultType) instanceof PyCallableParameterVariadicType paramVariadic - ? Ref.create(paramVariadic) : null); - } - } - } - } - } - return null; - } - - private static @NotNull PyTypeVarType.Variance getTypeVarVarianceFromDeclaration(@NotNull PyCallExpression assignedCall) { - boolean covariant = PyEvaluator.evaluateAsBooleanNoResolve(assignedCall.getKeywordArgument("covariant"), false); - boolean contravariant = PyEvaluator.evaluateAsBooleanNoResolve(assignedCall.getKeywordArgument("contravariant"), false); - boolean inferVariance = PyEvaluator.evaluateAsBooleanNoResolve(assignedCall.getKeywordArgument("infer_variance"), false); - - if (covariant && !contravariant) { - return PyTypeVarType.Variance.COVARIANT; - } - else if (contravariant && !covariant) { - return PyTypeVarType.Variance.CONTRAVARIANT; - } - else if (inferVariance) { - return PyTypeVarType.Variance.INFER_VARIANCE; - } - else { - return PyTypeVarType.Variance.INVARIANT; - } - } - - @ApiStatus.Internal - public static @Nullable PyAstTypeParameter.Kind getTypeParameterKindFromDeclaration(@NotNull PyCallExpression callExpression, - @NotNull TypeEvalContext context) { - final PyExpression callee = callExpression.getCallee(); - if (callee != null) { - final Collection calleeQNames = resolveToQualifiedNames(callee, context); - if (calleeQNames.contains(TYPE_VAR_TUPLE) || calleeQNames.contains(TYPE_VAR_TUPLE_EXT)) return PyAstTypeParameter.Kind.TypeVarTuple; - if (calleeQNames.contains(TYPE_VAR) || calleeQNames.contains(TYPE_VAR_EXT)) return PyAstTypeParameter.Kind.TypeVar; - if (calleeQNames.contains(PARAM_SPEC) || calleeQNames.contains(PARAM_SPEC_EXT)) return PyAstTypeParameter.Kind.ParamSpec; - } - return null; - } - - @ApiStatus.Internal - public static @Nullable PyTypeParameterType getTypeParameterTypeFromTypeParameter(@NotNull PyTypeParameter typeParameter, - @NotNull TypeEvalContext context) { - return staticWithCustomContext(context, c -> getTypeParameterTypeFromTypeParameter(typeParameter, c)); - } - - private static @Nullable PyTypeParameterType getTypeParameterTypeFromTypeParameter(@NotNull PsiElement element, - @NotNull Context context) { - if (element instanceof PyTypeParameter typeParameter) { - String name = typeParameter.getName(); - if (name == null) { - return null; - } - - ScopeOwner typeParameterOwner = ScopeUtil.getScopeOwner(typeParameter); - PyQualifiedNameOwner scopeOwner = typeParameterOwner instanceof PyQualifiedNameOwner qualifiedNameOwner ? qualifiedNameOwner : null; - - String defaultExpressionText = typeParameter.getDefaultExpressionText(); - PyExpression defaultExpression = defaultExpressionText != null - ? PyUtil.createExpressionFromFragment(defaultExpressionText, typeParameter) - : null; - Ref defaultType = null; - if (defaultExpression != null) { - final PyExpression defaultExprWithoutParens = PyPsiUtils.flattenParens(defaultExpression); - defaultType = defaultExprWithoutParens != null - ? getTypePreventingRecursion(defaultExprWithoutParens, context) - : Ref.create(); - } - - PyQualifiedNameOwner declarationElement = as(element, PyQualifiedNameOwner.class); - - switch (typeParameter.getKind()) { - case TypeVar -> { - List<@Nullable PyType> constraints = List.of(); - PyType boundType = null; - String boundExpressionText = typeParameter.getBoundExpressionText(); - PyExpression boundExpression = boundExpressionText != null - ? PyPsiUtils.flattenParens(PyUtil.createExpressionFromFragment(boundExpressionText, typeParameter)) - : null; - if (boundExpression instanceof PyTupleExpression tupleExpression) { - constraints = ContainerUtil.map(tupleExpression.getElements(), expr -> Ref.deref(getTypePreventingRecursion(expr, context))); - } - else if (boundExpression != null) { - boundType = Ref.deref(getTypePreventingRecursion(boundExpression, context)); - } - return new PyTypeVarTypeImpl(name, constraints, boundType, defaultType, PyTypeVarType.Variance.INFER_VARIANCE) - .withScopeOwner(scopeOwner) - .withDeclarationElement(declarationElement); - } - case ParamSpec -> { - return new PyParamSpecType(name) - .withScopeOwner(scopeOwner) - .withDefaultType( - Ref.deref(defaultType) instanceof PyCallableParameterVariadicType variadicType ? Ref.create(variadicType) : null) - .withDeclarationElement(declarationElement); - } - case TypeVarTuple -> { - return new PyTypeVarTupleTypeImpl(name) - .withScopeOwner(scopeOwner) - .withDefaultType(Ref.deref(defaultType) instanceof PyPositionalVariadicType variadicType ? Ref.create(variadicType) : null) - .withDeclarationElement(declarationElement); - } - } - } - return null; - } - - private static @Nullable Ref getTypePreventingRecursion(@NotNull PyExpression expression, @NotNull Context context) { - return doPreventingRecursion(expression, false, () -> getType(expression, context)); - } - - // See https://peps.python.org/pep-0484/#scoping-rules-for-type-variables - private static @Nullable PyQualifiedNameOwner getTypeParameterScope(@NotNull String name, - @NotNull PyExpression typeHint, - @NotNull Context context) { - if (!context.isComputeTypeParameterScopeEnabled()) return null; - - PsiElement typeHintContext = getStubRetainedTypeHintContext(typeHint); - List typeParamOwnerCandidates = - StreamEx.iterate(typeHintContext, Objects::nonNull, owner -> PsiTreeUtil.getStubOrPsiParentOfType(owner, ScopeOwner.class)) - .filter(owner -> owner instanceof PyFunction || owner instanceof PyClass) - .select(PyQualifiedNameOwner.class) - .toList(); - - PyQualifiedNameOwner closestOwner = ContainerUtil.getFirstItem(typeParamOwnerCandidates); - if (closestOwner instanceof PyFunction) { - Optional typeParameterType = StreamEx.of(typeParamOwnerCandidates) - .skip(1) - .map(owner -> findSameTypeParameterInDefinition(owner, name, context)) - .nonNull() - .findFirst(); - if (typeParameterType.isPresent()) { - return typeParameterType.get().getScopeOwner(); - } - } - if (closestOwner != null) { - boolean prevComputeTypeParameterScope = context.setComputeTypeParameterScopeEnabled(false); - try { - return findSameTypeParameterInDefinition(closestOwner, name, context) != null ? closestOwner : null; - } - finally { - context.setComputeTypeParameterScopeEnabled(prevComputeTypeParameterScope); - } - } - - // old-style type aliases of form `ListOf = list[T]` or `ListOf: TypeAlias = list[T]` - PyAssignmentStatement assignment = PsiTreeUtil.getParentOfType(typeHintContext, PyAssignmentStatement.class, false, PyStatement.class); - if (assignment != null) { - PyExpression assignedValue = PyPsiUtils.flattenParens(assignment.getAssignedValue()); - if (PsiTreeUtil.isAncestor(assignedValue, typeHintContext, false)) { - if (PyPsiUtils.flattenParens(assignment.getLeftHandSideExpression()) instanceof PyTargetExpression target) { - boolean isTypeParamDeclaration = assignedValue instanceof PyCallExpression callExpr && - getTypeParameterKindFromDeclaration(callExpr, context.getTypeContext()) != null; - - if (!isTypeParamDeclaration && PyTypingAliasStubType.looksLikeTypeHint(assignedValue)) { - return target; - } - } - } - } - - return null; - } - - private static @Nullable PyTypeParameterType findSameTypeParameterInDefinition(@NotNull PyQualifiedNameOwner owner, - @NotNull String name, - @NotNull Context context) { - // At this moment, the definition of the TypeVar should be the type alias at the top of the stack. - // While evaluating type hints of enclosing functions' parameters, resolving to the same TypeVar - // definition shouldn't trigger the protection against recursive aliases, so we manually remove - // it from the top for the time being. - if (context.isTypeAliasStackEmpty()) { - return null; - } - PyQualifiedNameOwner typeVarDeclaration = context.popTypeAlias(); - assert typeVarDeclaration instanceof PyTargetExpression; - try { - final Iterable typeParameters; - if (owner instanceof PyClass cls) { - typeParameters = collectTypeParameters(cls, context); - } - else if (owner instanceof PyFunction function) { - typeParameters = collectTypeParameters(function, context.getTypeContext()); - } - else { - typeParameters = List.of(); - } - return ContainerUtil.find(typeParameters, type -> name.equals(type.getName())); - } - finally { - context.pushTypeAlias(typeVarDeclaration); - } - } - - @ApiStatus.Internal - public static @NotNull Iterable collectTypeParameters(@NotNull PyFunction function, - @NotNull TypeEvalContext context) { - return StreamEx.of(function.getParameterList().getParameters()) - .select(PyNamedParameter.class) - .map(parameter -> new PyTypingTypeProvider().getParameterType(parameter, function, context)) - .append(new PyTypingTypeProvider().getReturnType(function, context)) - .map(Ref::deref) - .map(paramType -> PyTypeChecker.collectGenerics(paramType, context)) - .flatMap(generics -> StreamEx.of(generics.getTypeVars()) - .append(generics.getParamSpecs()) - .append(generics.getTypeVarTuples()) - ); - } - - private static @NotNull PsiElement getStubRetainedTypeHintContext(@NotNull PsiElement typeHintExpression) { - PsiFile containingFile = typeHintExpression.getContainingFile(); - // Values from PSI stubs and regular type comments - PsiElement fragmentOwner = containingFile.getContext(); - if (fragmentOwner != null) { - return fragmentOwner; - } - // Values from function type comments and string literals - else if (containingFile instanceof PyFunctionTypeAnnotationFile || containingFile instanceof PyTypeHintFile) { - return PyPsiUtils.getRealContext(typeHintExpression); - } - else { - return typeHintExpression; - } - } - - public static @Nullable PyPositionalVariadicType getUnpackedType(@NotNull PsiElement element, @NotNull TypeEvalContext context) { - Ref<@Nullable PyType> typeRef = getTypeFromStarExpression(element, context); - if (typeRef == null) { - typeRef = getTypeFromUnpackOperator(element, context); - } - if (typeRef == null) { - return null; - } - var expressionType = typeRef.get(); - if (expressionType instanceof PyTupleType tupleType) { - return new PyUnpackedTupleTypeImpl(tupleType.getElementTypes(), tupleType.isHomogeneous()); - } - if (expressionType instanceof PyTypeVarTupleType typeVarTupleType) { - return typeVarTupleType; - } - return null; - } - - private static @Nullable Ref<@Nullable PyType> getTypeFromUnpackOperator(@NotNull PsiElement element, @NotNull TypeEvalContext context) { - if (!(element instanceof PySubscriptionExpression subscriptionExpr) || - !resolvesToQualifiedNames(subscriptionExpr.getOperand(), context, UNPACK, UNPACK_EXT)) { - return null; - } - PyExpression indexExpression = subscriptionExpr.getIndexExpression(); - if (!(indexExpression instanceof PyReferenceExpression || indexExpression instanceof PySubscriptionExpression)) return null; - return Ref.create(Ref.deref(getType(indexExpression, context))); - } - - private static @Nullable Ref<@Nullable PyType> getTypeFromStarExpression(@NotNull PsiElement element, @NotNull TypeEvalContext context) { - if (!(element instanceof PyStarExpression starExpression)) return null; - PyExpression starredExpression = starExpression.getExpression(); - if (!(starredExpression instanceof PyReferenceExpression || starredExpression instanceof PySubscriptionExpression)) return null; - return Ref.create(Ref.deref(getType(starredExpression, context))); - } - - private static @NotNull List getIndexTypes(@NotNull PySubscriptionExpression expression, @NotNull Context context) { - final List types = new ArrayList<>(); - final PyExpression indexExpr = expression.getIndexExpression(); - if (indexExpr instanceof PyTupleExpression tupleExpr) { - for (PyExpression expr : tupleExpr.getElements()) { - types.add(Ref.deref(getType(expr, context))); - } - } - else if (indexExpr != null) { - types.add(Ref.deref(getType(indexExpr, context))); - } - return types; - } - - private static @Nullable PyCollectionType parameterizeClassDefaultAware(@NotNull PyClass pyClass, - @NotNull List actualTypeParams, - @NotNull Context context) { - PyCollectionType genericDefinitionType = - doPreventingRecursion(pyClass, false, () -> PyTypeChecker.findGenericDefinitionType(pyClass, context.getTypeContext())); - if (genericDefinitionType != null && ContainerUtil.exists(genericDefinitionType.getElementTypes(), - t -> t instanceof PyTypeParameterType typeParameterType && - typeParameterType.getDefaultType() != null)) { - - PyType parameterizedType = PyTypeChecker.parameterizeType(genericDefinitionType, actualTypeParams, context.myContext); - if (parameterizedType instanceof PyCollectionType collectionType) { - return collectionType; - } - } - return null; - } - - private static @Nullable Ref getTypeFromTypeAlias(@NotNull PyQualifiedNameOwner alias, - @NotNull PsiElement typeHint, - @NotNull PsiElement element, - @NotNull Context context) { - if (element instanceof PyExpression assignedExpression) { - if (alias instanceof PyTypeAliasStatement typeAliasStatement) { - return getTypeFromTypeAliasStatement(typeAliasStatement, typeHint, assignedExpression, context); - } - - @Nullable Ref assignedTypeRef = getType(assignedExpression, context); - if (assignedTypeRef != null) { - @Nullable PyType assignedType = assignedTypeRef.get(); - if (assignedType == null) { - return assignedTypeRef; - } - if (typeHint instanceof PySubscriptionExpression subscriptionExpr) { - List indexTypes = getIndexTypes(subscriptionExpr, context); - return Ref.create(PyTypeChecker.parameterizeType(assignedType, indexTypes, context.myContext)); - } - if (typeHint instanceof PyReferenceExpression) { - if (!(assignedType instanceof PyTypeParameterType)) { - List typeAliasTypeParams = - PyTypeChecker.collectGenerics(assignedType, context.getTypeContext()).getAllTypeParameters(); - if (!typeAliasTypeParams.isEmpty()) { - return Ref.create(PyTypeChecker.parameterizeType(assignedType, List.of(), context.myContext)); - } - return Ref.create(assignedType); - } - } - } - } - return null; - } - - private static @Nullable Ref getTypeFromTypeAliasStatement(@NotNull PyTypeAliasStatement typeAliasStatement, - @NotNull PsiElement typeHint, - @NotNull PyExpression assignedExpression, - @NotNull Context context) { - @Nullable Ref assignedTypeRef = getType(assignedExpression, context); - if (assignedTypeRef != null) { - PyType assignedType = assignedTypeRef.get(); - if (assignedType == null) { - return assignedTypeRef; - } - List indexTypes = typeHint instanceof PySubscriptionExpression subscriptionExpr - ? getIndexTypes(subscriptionExpr, context) - : Collections.emptyList(); - - List typeAliasTypeParams = collectTypeParametersFromTypeAliasStatement(typeAliasStatement, context); - if (!typeAliasTypeParams.isEmpty()) { - PyTypeChecker.GenericSubstitutions substitutions = - PyTypeChecker.mapTypeParametersToSubstitutions(typeAliasTypeParams, - indexTypes, - Option.USE_DEFAULTS, - Option.MAP_UNMATCHED_EXPECTED_TYPES_TO_ANY); - - return substitutions != null ? Ref.create(PyTypeChecker.substitute(assignedType, substitutions, context.myContext)) : null; - } - return assignedTypeRef; - } - return null; - } - - private static @Nullable PyType getParameterizedType(@NotNull PsiElement element, @NotNull Context context) { - if (element instanceof PySubscriptionExpression subscriptionExpr) { - final PyExpression operand = subscriptionExpr.getOperand(); - final PyExpression indexExpr = subscriptionExpr.getIndexExpression(); - if (indexExpr != null) { - final PyType operandType = Ref.deref(getType(operand, context)); - final List indexTypes = getIndexTypes(subscriptionExpr, context); - if (operandType != null) { - if (operandType instanceof PyClassType classType) { - if (!(operandType instanceof PyTupleType) && PyNames.TUPLE.equals(classType.getPyClass().getQualifiedName())) { - if (indexExpr instanceof PyTupleExpression) { - final PyExpression[] elements = ((PyTupleExpression)indexExpr).getElements(); - if (elements.length == 2 && elements[1] instanceof PyEllipsisLiteralExpression) { - return PyTupleType.createHomogeneous(element, indexTypes.get(0)); - } - } - return PyTupleType.create(element, indexTypes); - } - - if (isGeneric(classType, context.myContext)) { - PyCollectionType parameterizedType = parameterizeClassDefaultAware(classType.getPyClass(), indexTypes, context); - if (parameterizedType != null) { - return parameterizedType.toInstance(); - } - } - else { - return null; - } - } - return PyTypeChecker.parameterizeType(operandType, indexTypes, context.getTypeContext()); - } - } - } - return null; - } - - private static @Nullable PyType getCollection(@NotNull PsiElement element, @NotNull TypeEvalContext context) { - final String typingName = getQualifiedName(element); - - final String builtinName = BUILTIN_COLLECTION_CLASSES.get(typingName); - if (builtinName != null) return PyTypeParser.getTypeByName(element, builtinName, context); - - final String collectionName = COLLECTIONS_CLASSES.get(typingName); - if (collectionName != null) return PyTypeParser.getTypeByName(element, collectionName, context); - - return null; - } - - private static @NotNull List tryResolving(@NotNull PyExpression expression, @NotNull TypeEvalContext context) { - return ContainerUtil.map(tryResolvingWithAliases(expression, context), x -> x.getSecond()); - } - - private static @NotNull List> tryResolvingWithAliases(@NotNull PyExpression expression, - @NotNull TypeEvalContext context) { - final List> elements = new ArrayList<>(); - if (expression instanceof PyReferenceExpression) { - final List results; - if (context.maySwitchToAST(expression)) { - final PyResolveContext resolveContext = PyResolveContext.defaultContext(context); - results = PyUtil.multiResolveTopPriority(expression, resolveContext); - } - else { - results = tryResolvingOnStubs((PyReferenceExpression)expression, context); - } - for (PsiElement element : results) { - final PyClass cls = PyUtil.turnConstructorIntoClass(as(element, PyFunction.class)); - if (cls != null) { - elements.add(Pair.create(null, cls)); - continue; - } - final String name = element != null ? getQualifiedName(element) : null; - if (name != null && OPAQUE_NAMES.contains(name)) { - elements.add(Pair.create(null, element)); - continue; - } - // Presumably, a TypeVar definition or a type alias - if (element instanceof PyTargetExpression targetExpr) { - final PyExpression assignedValue = PyTypingAliasStubType.getAssignedValueStubLike(targetExpr); - if (assignedValue != null) { - elements.add(Pair.create(targetExpr, assignedValue)); - continue; - } - } - if (element instanceof PyTypeAliasStatement typeAliasStatement) { - PyExpression assignedValue; - if (context.maySwitchToAST(typeAliasStatement)) { - assignedValue = typeAliasStatement.getTypeExpression(); - } - else { - String assignedTypeText = typeAliasStatement.getTypeExpressionText(); - assignedValue = assignedTypeText != null ? PyUtil.createExpressionFromFragment(assignedTypeText, typeAliasStatement) : null; - } - if (assignedValue != null) { - elements.add(Pair.create(typeAliasStatement, assignedValue)); - continue; - } - } - if (element != null) { - elements.add(Pair.create(null, element)); - } - } - } - if (expression instanceof PySubscriptionExpression subscriptionExpr) { - // Possibly a parameterized type alias - PyExpression operandExpression = subscriptionExpr.getOperand(); - List> results = tryResolvingWithAliases(operandExpression, context); - for (Pair pair : results) { - // If the parameterized type is a type alias - if (pair.getFirst() != null && pair.getSecond() != null) { - elements.add(Pair.create(pair.getFirst(), pair.getSecond())); - } - } - } - return !elements.isEmpty() ? elements : Collections.singletonList(Pair.create(null, expression)); - } - - private static @NotNull List tryResolvingOnStubs(@NotNull PyReferenceExpression expression, - @NotNull TypeEvalContext context) { - - final QualifiedName qualifiedName = expression.asQualifiedName(); - final PyFile pyFile = as(FileContextUtil.getContextFile(expression), PyFile.class); - - PsiElement anchor = expression.getContainingFile().getContext(); - ScopeOwner scopeOwner; - - if (anchor == null) { - scopeOwner = pyFile; - } - else if (anchor instanceof ScopeOwner anchorAsScope) { - scopeOwner = anchorAsScope; - } - else { - scopeOwner = ScopeUtil.getScopeOwner(anchor); - } - - if (scopeOwner != null && qualifiedName != null) { - return PyResolveUtil.resolveQualifiedNameInScope(qualifiedName, scopeOwner, context); - } - return Collections.singletonList(expression); - } - - public static @NotNull Collection resolveToQualifiedNames(@NotNull PyExpression expression, @NotNull TypeEvalContext context) { - final Set names = new LinkedHashSet<>(); - for (PsiElement resolved : tryResolving(expression, context)) { - final String name = getQualifiedName(resolved); - if (name != null) { - names.add(name); - } - } - return names; - } - - private static @Nullable String getQualifiedName(@NotNull PsiElement element) { - if (element instanceof PyQualifiedNameOwner qualifiedNameOwner) { - return qualifiedNameOwner.getQualifiedName(); - } - return null; - } - - public static @Nullable PyType toAsyncIfNeeded(@NotNull PyFunction function, @Nullable PyType returnType) { - if (function.isAsync() && function.isAsyncAllowed()) { - if (!function.isGenerator()) { - return wrapInCoroutineType(returnType, function); - } - var desc = GeneratorTypeDescriptor.fromGenerator(returnType); - if (desc != null) { - final PyClass classType = PyPsiFacade.getInstance(function.getProject()).createClassByQName(ASYNC_GENERATOR, function); - final List generics = Arrays.asList(desc.yieldType, desc.sendType); - return classType != null ? new PyCollectionTypeImpl(classType, false, generics) : null; - } - } - return returnType; - } - - /** - * Bound narrowed types shouldn't leak out of its scope, since it is bound to a particular call site. - */ - public static @Nullable PyType removeNarrowedTypeIfNeeded(@Nullable PyType type) { - if (type instanceof PyNarrowedType pyNarrowedType && pyNarrowedType.isBound()) { - return PyBuiltinCache.getInstance(pyNarrowedType.getOriginal()).getBoolType(); - } - else { - return type; - } - } - - private static @Nullable PyType wrapInCoroutineType(@Nullable PyType returnType, @NotNull PsiElement resolveAnchor) { - final PyClass coroutine = PyPsiFacade.getInstance(resolveAnchor.getProject()).createClassByQName(COROUTINE, resolveAnchor); - return coroutine != null ? new PyCollectionTypeImpl(coroutine, false, Arrays.asList(null, null, returnType)) : null; - } - - public static @Nullable PyType wrapInGeneratorType(@Nullable PyType elementType, - @Nullable PyType sendType, - @Nullable PyType returnType, - @NotNull PsiElement anchor) { - final PyClass generator = PyPsiFacade.getInstance(anchor.getProject()).createClassByQName(GENERATOR, anchor); - return generator != null ? new PyCollectionTypeImpl(generator, false, Arrays.asList(elementType, sendType, returnType)) : null; - } - - public record GeneratorTypeDescriptor( - @Nullable PyType yieldType, - @Nullable PyType sendType, - @Nullable PyType returnType, - boolean isAsync - ) { - /** - * Extracts type parameters from typing.Generator and typing.AsyncGenerator - */ - public static @Nullable GeneratorTypeDescriptor fromGenerator(@Nullable PyType type) { - if (!(type instanceof PyClassType classType)) return null; - - final String qName = classType.getClassQName(); - if (qName == null) return null; - - boolean isAsync = ASYNC_GENERATOR.equals(qName); - if (!isAsync && !GENERATOR.equals(qName)) return null; - - final PyType noneType = PyBuiltinCache.getInstance(classType.getPyClass()).getNoneType(); - - PyType yieldType = null; - PyType sendType = noneType; - PyType returnType = isAsync ? null : noneType; - if (type instanceof PyCollectionType genericType) { - yieldType = ContainerUtil.getOrElse(genericType.getElementTypes(), 0, yieldType); - sendType = ContainerUtil.getOrElse(genericType.getElementTypes(), 1, sendType); - returnType = ContainerUtil.getOrElse(genericType.getElementTypes(), 2, returnType); - } - return new GeneratorTypeDescriptor(yieldType, sendType, returnType, isAsync); - } - - /** - * Unlike {@link #fromGenerator}, this method can also extract yield type from Protocol types like typing.Iterable - */ - public static @Nullable GeneratorTypeDescriptor fromGeneratorOrProtocol(@Nullable PyType type, @NotNull TypeEvalContext context) { - if (!(type instanceof PyClassType classType)) return null; - - GeneratorTypeDescriptor desc = fromGenerator(type); - if (desc != null) { - return desc; - } - - if (PyProtocolsKt.isProtocol(classType, context)) { - PyType yieldType; - - PyType syncUpcast = PyTypeUtil.convertToType(classType, "typing.Iterable", classType.getPyClass(), context); - if (syncUpcast instanceof PyCollectionType collectionType) { - yieldType = collectionType.getIteratedItemType(); - return new GeneratorTypeDescriptor(yieldType, null, null, false); - } - PyType asyncUpcast = PyTypeUtil.convertToType(classType, "typing.AsyncIterable", classType.getPyClass(), context); - if (asyncUpcast instanceof PyCollectionType asyncCollectionType) { - yieldType = asyncCollectionType.getIteratedItemType(); - return new GeneratorTypeDescriptor(yieldType, null, null, true); - } - - // Here we try to understand a yield type by return type of __next__ method of protocol specified in annotation. - // We cannot use convertToType with typing.Iterator here, as it inherits from typing.Iterable - // and requires both __iter__ and __next__, while it should be possible to decide the yield type only by __next__. - // TODO: unify logic with PyTargetExpressionImpl.getIterationType (PY-82453) - PyFunction next = classType.getPyClass().findMethodByName(PyNames.DUNDER_NEXT, true, context); - if (next != null) { - yieldType = context.getReturnType(next); - yieldType = PyTypeChecker.substitute(yieldType, PyTypeChecker.unifyReceiver(classType, context), context); - return new GeneratorTypeDescriptor(yieldType, null, null, false); - } - - PyFunction anext = classType.getPyClass().findMethodByName(PyNames.ANEXT, true, context); - if (anext != null) { - yieldType = Ref.deref(unwrapCoroutineReturnType(context.getReturnType(anext))); - yieldType = PyTypeChecker.substitute(yieldType, PyTypeChecker.unifyReceiver(classType, context), context); - return new GeneratorTypeDescriptor(yieldType, null, null, true); - } - } - return null; - } - } - - public static @Nullable Ref unwrapCoroutineReturnType(@Nullable PyType coroutineType) { - final PyCollectionType genericType = as(coroutineType, PyCollectionType.class); - - if (genericType != null) { - var qName = genericType.getClassQName(); - - if (AWAITABLE.equals(qName)) { - return Ref.create(ContainerUtil.getOrElse(genericType.getElementTypes(), 0, null)); - } - - if (COROUTINE.equals(qName)) { - return Ref.create(ContainerUtil.getOrElse(genericType.getElementTypes(), 2, null)); - } - } - - return null; - } - - public static @Nullable Ref coroutineOrGeneratorElementType(@Nullable PyType coroutineOrGeneratorType) { - final PyCollectionType genericType = as(coroutineOrGeneratorType, PyCollectionType.class); - final PyClassType classType = as(coroutineOrGeneratorType, PyClassType.class); - - if (genericType != null && classType != null) { - var qName = classType.getClassQName(); - - if (AWAITABLE.equals(qName)) { - return Ref.create(ContainerUtil.getOrElse(genericType.getElementTypes(), 0, null)); - } - - if (ArrayUtil.contains(qName, COROUTINE, GENERATOR)) { - return Ref.create(ContainerUtil.getOrElse(genericType.getElementTypes(), 2, null)); - } - } - - return null; - } - - /** - * Checks whether the given assignment is type hinted with {@code typing.TypeAlias}. - *

- * It can be done either with a variable annotation or a type comment. - */ - public static boolean isExplicitTypeAlias(@NotNull PyAssignmentStatement assignment, @NotNull TypeEvalContext context) { - PyTargetExpression target = as(ArrayUtil.getFirstElement(assignment.getTargets()), PyTargetExpression.class); - if (target == null) { - return false; - } - return isExplicitTypeAlias(target, context); - } - - public static boolean isExplicitTypeAlias(@NotNull PyTargetExpression targetExpression, @NotNull TypeEvalContext context) { - PyExpression annotationValue = getAnnotationValue(targetExpression, context); - if (annotationValue instanceof PyReferenceExpression) { - return resolvesToQualifiedNames(annotationValue, context, TYPE_ALIAS, TYPE_ALIAS_EXT); - } - String typeCommentAnnotation = targetExpression.getTypeCommentAnnotation(); - if (typeCommentAnnotation != null) { - PyExpression commentValue = PyUtil.createExpressionFromFragment(typeCommentAnnotation, targetExpression); - if (commentValue instanceof PyReferenceExpression) { - return resolvesToQualifiedNames(commentValue, context, TYPE_ALIAS, TYPE_ALIAS_EXT); - } - } - return false; - } - - /** - * Detects whether the given element belongs to a self-evident type hint. Namely, these are: - *

    - *
  • function and variable annotations
    • - *
  • type comments
    • - *
  • explicit type aliases marked with {@code TypeAlias}
    • - *
- * Note that {@code element} can belong to their AST directly or be a part of an injection inside one of such elements. - */ - public static boolean isInsideTypeHint(@NotNull PsiElement element, @NotNull TypeEvalContext context) { - final PsiElement realContext = PyPsiUtils.getRealContext(element); - - if (PsiTreeUtil.getParentOfType(realContext, PyAnnotation.class, false, PyStatement.class) != null) { - return true; - } - - final PsiComment comment = PsiTreeUtil.getParentOfType(realContext, PsiComment.class, false, PyStatement.class); - if (comment != null && getTypeCommentValue(comment.getText()) != null) { - return true; - } - - PyAssignmentStatement assignment = PsiTreeUtil.getParentOfType(realContext, PyAssignmentStatement.class, false, PyStatement.class); - if (assignment != null && - PsiTreeUtil.isAncestor(assignment.getAssignedValue(), realContext, false) && - isExplicitTypeAlias(assignment, context)) { - return true; - } - - PyTypeAliasStatement typeAlias = PsiTreeUtil.getParentOfType(realContext, PyTypeAliasStatement.class, false, PyStatement.class); - if (typeAlias != null && PsiTreeUtil.isAncestor(typeAlias.getTypeExpression(), realContext, false)) { - return true; - } - - return false; - } - - @Override - protected T withCustomContext(@NotNull TypeEvalContext context, @NotNull Function<@NotNull Context, T> delegate) { - return staticWithCustomContext(context, delegate); - } - - private static T staticWithCustomContext(@NotNull TypeEvalContext context, @NotNull Function<@NotNull Context, T> delegate) { - return staticWithCustomContext(context, false, delegate); - } - - private static T staticWithCustomContext(@NotNull TypeEvalContext context, - boolean useFqn, - @NotNull Function<@NotNull Context, T> delegate) { - Context customContext = context.getProcessingContext().get(TYPE_HINT_EVAL_CONTEXT); - boolean firstEntrance = customContext == null; - if (firstEntrance) { - customContext = new Context(context, useFqn); - context.getProcessingContext().put(TYPE_HINT_EVAL_CONTEXT, customContext); - } - try { - return delegate.apply(customContext); - } - finally { - if (firstEntrance) { - context.getProcessingContext().put(TYPE_HINT_EVAL_CONTEXT, null); - } - } - } - - @ApiStatus.Internal - public static final class Context { - private final @NotNull TypeEvalContext myContext; - private final @NotNull Stack myTypeAliasStack = new Stack<>(); - private final @NotNull Set myClassSet = new HashSet<>(); - private boolean myComputeTypeParameterScope = true; - private final boolean myUseFqn; - - private Context(@NotNull TypeEvalContext context) { - myContext = context; - myUseFqn = false; - recomputeStrongHashValue(); - } - - private Context(@NotNull TypeEvalContext context, boolean useFqn) { - myContext = context; - myUseFqn = useFqn; - recomputeStrongHashValue(); - } - - public @NotNull TypeEvalContext getTypeContext() { - return myContext; - } - - public @NotNull Stack getTypeAliasStack() { - return myTypeAliasStack; + class Context(val typeContext: TypeEvalContext, val typeRepresentationMode: Boolean = false) { + val typeAliasStack: Stack = Stack() + private val myClassSet: MutableSet = HashSet() + var isComputeTypeParameterScopeEnabled: Boolean = true + private set + + init { + recomputeStrongHashValue() } // Explicit API for manipulating type alias stack - public boolean containsTypeAlias(@NotNull PyQualifiedNameOwner alias) { - return myTypeAliasStack.contains(alias); + fun containsTypeAlias(alias: PyQualifiedNameOwner): Boolean { + return alias in typeAliasStack } - public void addTypeAlias(@NotNull PyQualifiedNameOwner alias) { - myTypeAliasStack.add(alias); - recomputeStrongHashValue(); + fun addTypeAlias(alias: PyQualifiedNameOwner) { + typeAliasStack.add(alias) + recomputeStrongHashValue() } - public void removeTypeAlias(@NotNull PyQualifiedNameOwner alias) { - myTypeAliasStack.remove(alias); - recomputeStrongHashValue(); + fun removeTypeAlias(alias: PyQualifiedNameOwner) { + typeAliasStack.remove(alias) + recomputeStrongHashValue() } - public boolean isTypeAliasStackEmpty() { - return myTypeAliasStack.isEmpty(); + val isTypeAliasStackEmpty: Boolean + get() = typeAliasStack.isEmpty() + + fun peekTypeAlias(): PyQualifiedNameOwner? { + return if (typeAliasStack.isEmpty()) null else typeAliasStack.peek() } - public @Nullable PyQualifiedNameOwner peekTypeAlias() { - return myTypeAliasStack.isEmpty() ? null : myTypeAliasStack.peek(); + fun pushTypeAlias(alias: PyQualifiedNameOwner) { + typeAliasStack.push(alias) + recomputeStrongHashValue() } - public void pushTypeAlias(@NotNull PyQualifiedNameOwner alias) { - myTypeAliasStack.push(alias); - recomputeStrongHashValue(); + fun popTypeAlias(): PyQualifiedNameOwner? { + val res = if (typeAliasStack.isEmpty()) null else typeAliasStack.pop() + recomputeStrongHashValue() + return res } - public @Nullable PyQualifiedNameOwner popTypeAlias() { - PyQualifiedNameOwner res = myTypeAliasStack.isEmpty() ? null : myTypeAliasStack.pop(); - recomputeStrongHashValue(); - return res; + fun addClassDeclaration(pyClass: PyClass): Boolean { + return myClassSet.add(pyClass) } - public boolean addClassDeclaration(@NotNull PyClass pyClass) { - return myClassSet.add(pyClass); + fun removeClassDeclaration(pyClass: PyClass) { + myClassSet.remove(pyClass) } - public void removeClassDeclaration(@NotNull PyClass pyClass) { - myClassSet.remove(pyClass); + fun setComputeTypeParameterScopeEnabled(value: Boolean): Boolean { + val prev = isComputeTypeParameterScopeEnabled + isComputeTypeParameterScopeEnabled = value + recomputeStrongHashValue() + return prev } - public boolean isComputeTypeParameterScopeEnabled() { - return myComputeTypeParameterScope; + fun getKnownType(expression: PyExpression): PyType? { + return typeContext.getContextTypeCache()[expression to contextStrongHashValue] } - public boolean setComputeTypeParameterScopeEnabled(boolean value) { - boolean prev = myComputeTypeParameterScope; - myComputeTypeParameterScope = value; - recomputeStrongHashValue(); - return prev; + fun assumeType(expression: PyExpression, type: PyType) { + typeContext.getContextTypeCache()[expression to contextStrongHashValue] = type } - public @Nullable PyType getKnownType(@NotNull PyExpression expression) { - //noinspection SuspiciousMethodCalls - return myContext.getContextTypeCache().get(new kotlin.Pair<>(expression, getContextStrongHashValue())); + private var myContextStrongHashValue: HashValue128? = null + + private fun recomputeStrongHashValue() { + myContextStrongHashValue = Hashing.xxh3_128().hashCharsTo128Bits( + buildList { + add(if (isComputeTypeParameterScopeEnabled) "1" else "0") + add(typeAliasStack.map { it!!.qualifiedName }) + add(myClassSet.map { it!!.qualifiedName }) + }.joinToString("#") + ) } - public void assumeType(@NotNull PyExpression expression, @NotNull PyType type) { - myContext.getContextTypeCache().put(new kotlin.Pair<>(expression, getContextStrongHashValue()), type); + private val contextStrongHashValue: HashValue128 + get() = myContextStrongHashValue!! + + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (other == null || javaClass != other.javaClass) return false + val context = other as Context + return typeContext == context.typeContext } - private @NotNull HashValue128 myContextStrongHashValue; + override fun hashCode(): Int { + return Objects.hash(typeContext) + } + } - private void recomputeStrongHashValue() { - myContextStrongHashValue = xxh3_128().hashCharsTo128Bits( - StreamEx.of(myComputeTypeParameterScope ? "1" : "0") - .append(myTypeAliasStack.stream().map(it -> it.getQualifiedName())) - .append(myClassSet.stream().map(it -> it.getQualifiedName())) - .joining("#") - ); + companion object { + @NlsSafe + const val TYPING: @NlsSafe String = "typing" + + const val GENERATOR: String = "typing.Generator" + const val ASYNC_GENERATOR: String = "typing.AsyncGenerator" + const val COROUTINE: String = "typing.Coroutine" + const val AWAITABLE: String = "typing.Awaitable" + const val NAMEDTUPLE: String = "typing.NamedTuple" + const val TYPED_DICT: String = "typing.TypedDict" + const val TYPED_DICT_EXT: String = "typing_extensions.TypedDict" + const val TYPE_GUARD: String = "typing.TypeGuard" + const val TYPE_GUARD_EXT: String = "typing_extensions.TypeGuard" + const val TYPE_IS: String = "typing.TypeIs" + const val TYPE_IS_EXT: String = "typing_extensions.TypeIs" + const val GENERIC: String = "typing.Generic" + const val PROTOCOL: String = "typing.Protocol" + const val PROTOCOL_EXT: String = "typing_extensions.Protocol" + const val TYPE: String = "typing.Type" + const val ANY: String = "typing.Any" + const val NEW_TYPE: String = "typing.NewType" + const val CALLABLE: String = "typing.Callable" + const val CALLABLE_EXT: String = "typing_extensions.Callable" + const val MAPPING: String = "typing.Mapping" + const val MAPPING_GET: String = "typing.Mapping.get" + private const val LIST = "typing.List" + private const val DICT = "typing.Dict" + private const val DEFAULT_DICT = "typing.DefaultDict" + private const val ORDERED_DICT = "typing.OrderedDict" + private const val SET = "typing.Set" + private const val FROZEN_SET = "typing.FrozenSet" + private const val COUNTER = "typing.Counter" + private const val DEQUE = "typing.Deque" + private const val TUPLE = "typing.Tuple" + const val CLASS_VAR: String = "typing.ClassVar" + const val TYPE_VAR: String = "typing.TypeVar" + const val TYPE_VAR_EXT: String = "typing_extensions.TypeVar" + const val TYPE_VAR_TUPLE: String = "typing.TypeVarTuple" + const val TYPE_VAR_TUPLE_EXT: String = "typing_extensions.TypeVarTuple" + const val PARAM_SPEC: String = "typing.ParamSpec" + const val PARAM_SPEC_EXT: String = "typing_extensions.ParamSpec" + private const val CHAIN_MAP = "typing.ChainMap" + const val UNION: String = "typing.Union" + const val CONCATENATE: String = "typing.Concatenate" + const val CONCATENATE_EXT: String = "typing_extensions.Concatenate" + const val OPTIONAL: String = "typing.Optional" + const val NO_RETURN: String = "typing.NoReturn" + const val NEVER: String = "typing.Never" + const val NO_RETURN_EXT: String = "typing_extensions.NoReturn" + const val NEVER_EXT: String = "typing_extensions.Never" + const val FINAL: String = "typing.Final" + const val FINAL_EXT: String = "typing_extensions.Final" + const val LITERAL: String = "typing.Literal" + const val LITERAL_EXT: String = "typing_extensions.Literal" + const val LITERALSTRING: String = "typing.LiteralString" + const val LITERALSTRING_EXT: String = "typing_extensions.LiteralString" + const val ANNOTATED: String = "typing.Annotated" + const val ANNOTATED_EXT: String = "typing_extensions.Annotated" + const val TYPE_ALIAS: String = "typing.TypeAlias" + const val TYPE_ALIAS_EXT: String = "typing_extensions.TypeAlias" + const val TYPE_ALIAS_TYPE: String = "typing.TypeAliasType" + const val SPECIAL_FORM: String = "typing._SpecialForm" + const val SPECIAL_FORM_EXT: String = "typing_extensions._SpecialForm" + const val REQUIRED: String = "typing.Required" + const val REQUIRED_EXT: String = "typing_extensions.Required" + const val NOT_REQUIRED: String = "typing.NotRequired" + const val NOT_REQUIRED_EXT: String = "typing_extensions.NotRequired" + const val READONLY: String = "typing.ReadOnly" + const val READONLY_EXT: String = "typing_extensions.ReadOnly" + + val TYPE_PARAMETER_FACTORIES: Set = setOf( + TYPE_VAR, TYPE_VAR_EXT, + PARAM_SPEC, PARAM_SPEC_EXT, + TYPE_VAR_TUPLE, TYPE_VAR_TUPLE_EXT + ) + + val TYPE_DICT_QUALIFIERS: Set = + setOf(REQUIRED, REQUIRED_EXT, NOT_REQUIRED, NOT_REQUIRED_EXT, READONLY, READONLY_EXT) + + const val UNPACK: String = "typing.Unpack" + const val UNPACK_EXT: String = "typing_extensions.Unpack" + + const val SELF: String = "typing.Self" + const val SELF_EXT: String = "typing_extensions.Self" + + val TYPE_IGNORE_PATTERN: Pattern = Pattern.compile("#\\s*type:\\s*ignore\\s*(\\[[^]#]*])?($|(\\s.*))", Pattern.CASE_INSENSITIVE) + + const val ASSERT_TYPE: String = "typing.assert_type" + const val REVEAL_TYPE: String = "typing.reveal_type" + const val REVEAL_TYPE_EXT: String = "typing_extensions.reveal_type" + const val CAST: String = "typing.cast" + const val CAST_EXT: String = "typing_extensions.cast" + + val BUILTIN_COLLECTION_CLASSES: ImmutableMap = ImmutableMap.builder() + .put(LIST, "list") + .put(DICT, "dict") + .put(SET, PyNames.SET) + .put(FROZEN_SET, "frozenset") + .put(TUPLE, PyNames.TUPLE) + .build() + + private val COLLECTIONS_CLASSES = ImmutableMap.builder() + .put(DEFAULT_DICT, "collections.defaultdict") + .put(ORDERED_DICT, "collections.OrderedDict") + .put(COUNTER, "collections.Counter") + .put(DEQUE, "collections.deque") + .put(CHAIN_MAP, "collections.ChainMap") + .build() + + @JvmField + val TYPING_COLLECTION_CLASSES: ImmutableMap = ImmutableMap.builder() + .put("list", "List") + .put("dict", "Dict") + .put("set", "Set") + .put("frozenset", "FrozenSet") + .build() + + val TYPING_BUILTINS_GENERIC_ALIASES: ImmutableMap = ImmutableMap.builder() + .putAll(TYPING_COLLECTION_CLASSES.entries) + .put("type", "Type") + .put("tuple", "Tuple") + .build() + + @JvmField + val GENERIC_CLASSES: ImmutableSet = ImmutableSet.builder() // special forms + .add( + TUPLE, + GENERIC, + PROTOCOL, + CALLABLE, + CALLABLE_EXT, + TYPE, + CLASS_VAR, + FINAL, + LITERAL, + ANNOTATED, + REQUIRED, + NOT_REQUIRED, + READONLY + ) // type aliases + .add(UNION, OPTIONAL, LIST, DICT, DEFAULT_DICT, ORDERED_DICT, SET, FROZEN_SET, COUNTER, DEQUE, CHAIN_MAP) + .add(PROTOCOL_EXT, FINAL_EXT, LITERAL_EXT, ANNOTATED_EXT, REQUIRED_EXT, NOT_REQUIRED_EXT, READONLY_EXT) + .build() + + /** + * For the following names we shouldn't go further to the RHS of assignments, + * since they are not type aliases already and in typing.pyi are assigned to + * some synthetic values. + */ + val OPAQUE_NAMES: ImmutableSet = ImmutableSet.builder() + .add(PyKnownDecorator.TYPING_OVERLOAD.qualifiedName.toString()) + .add(ANY) + .add(TYPE_VAR) + .add(TYPE_VAR_EXT) + .add(TYPE_VAR_TUPLE) + .add(TYPE_VAR_TUPLE_EXT) + .add(GENERIC) + .add(PARAM_SPEC) + .add(PARAM_SPEC_EXT) + .add(CONCATENATE) + .add(CONCATENATE_EXT) + .add(TUPLE) + .add(CALLABLE) + .add(CALLABLE_EXT) + .add(TYPE) + .add(PyKnownDecorator.TYPING_NO_TYPE_CHECK.qualifiedName.toString()) + .add(PyKnownDecorator.TYPING_NO_TYPE_CHECK_EXT.qualifiedName.toString()) + .add(UNION) + .add(OPTIONAL) + .add(LIST) + .add(DICT) + .add(DEFAULT_DICT) + .add(ORDERED_DICT) + .add(SET) + .add(FROZEN_SET) + .add(PROTOCOL, PROTOCOL_EXT) + .add(CLASS_VAR) + .add(COUNTER) + .add(DEQUE) + .add(CHAIN_MAP) + .add(NO_RETURN, NO_RETURN_EXT) + .add(NEVER, NEVER_EXT) + .add(FINAL, FINAL_EXT) + .add(LITERAL, LITERAL_EXT) + .add(TYPED_DICT, TYPED_DICT_EXT) + .add(ANNOTATED, ANNOTATED_EXT) + .add(TYPE_ALIAS, TYPE_ALIAS_EXT) + .add(REQUIRED, REQUIRED_EXT) + .add(NOT_REQUIRED, NOT_REQUIRED_EXT) + .add(READONLY, READONLY_EXT) + .add(SELF, SELF_EXT) + .add(LITERALSTRING, LITERALSTRING_EXT) + .build() + + private val TYPE_HINT_EVAL_CONTEXT = Key.create("TYPE_HINT_EVAL_CONTEXT") + + private fun findParamTypeHintInFunctionTypeComment( + annotation: PyFunctionTypeAnnotation, + param: PyNamedParameter, + func: PyFunction, + ): PyExpression? { + val paramTypes = annotation.parameterTypeList.parameterTypes + if (paramTypes.size == 1 && paramTypes[0] is PyEllipsisLiteralExpression) { + return null + } + val startOffset = if (omitFirstParamInTypeComment(func, annotation)) 1 else 0 + val funcParams = listOf(*func.parameterList.parameters) + val i = funcParams.indexOf(param) - startOffset + if (i >= 0 && i < paramTypes.size) { + val paramTypeHint = paramTypes[i] + when (paramTypeHint) { + is PyStarExpression -> { + return paramTypeHint.expression + } + is PyDoubleStarExpression -> { + return paramTypeHint.expression + } + else -> { + return paramTypeHint + } + } + } + return null } - private @NotNull HashValue128 getContextStrongHashValue() { - return myContextStrongHashValue; + @JvmStatic + fun isGenerator(type: PyType): Boolean { + return type is PyCollectionType && GENERATOR == type.classQName } - @Override - public boolean equals(Object o) { - if (this == o) return true; - if (o == null || getClass() != o.getClass()) return false; - Context context = (Context)o; - return Objects.equals(myContext, context.myContext); + private fun createTypingGenericType(anchor: PsiElement): PyType { + return PyCustomType(GENERIC, null, false, true, getInstance(anchor).objectType) } - @Override - public int hashCode() { - return Objects.hash(myContext); + private fun createTypingProtocolType(anchor: PsiElement): PyType { + return PyCustomType(PROTOCOL, null, false, true, getInstance(anchor).objectType) + } + + fun createTypingCallableType(anchor: PsiElement): PyType { + return PyCustomType(CALLABLE, null, false, true, getInstance(anchor).objectType) + } + + private fun omitFirstParamInTypeComment(func: PyFunction, annotation: PyFunctionTypeAnnotation): Boolean { + return func.containingClass != null && func.modifier != PyAstFunction.Modifier.STATICMETHOD && annotation.parameterTypeList + .parameterTypes.size < func.parameterList.parameters.size + } + + @ApiStatus.Internal + @JvmStatic + fun getReturnTypeAnnotation(function: PyFunction, context: TypeEvalContext): PyExpression? { + val returnAnnotation: PyExpression? = getAnnotationValue(function, context) + if (returnAnnotation != null) { + return returnAnnotation + } + val functionAnnotation: PyFunctionTypeAnnotation? = getFunctionTypeAnnotation(function) + if (functionAnnotation != null) { + return functionAnnotation.returnType + } + return null + } + + fun getFunctionTypeAnnotation(function: PyFunction): PyFunctionTypeAnnotation? { + val comment = function.typeCommentAnnotation + if (comment == null) { + return null + } + val file = CachedValuesManager.getCachedValue( + function, + CachedValueProvider { + CachedValueProvider.Result.create( + PyFunctionTypeAnnotationFile( + function.typeCommentAnnotation!!, + function + ), function + ) + }) + return file.annotation + } + + private fun functionReturningCallSiteAsAType(function: PyFunction): Boolean { + val name = function.name + + if (PyNames.CLASS_GETITEM == name) return true + if (PyNames.GETITEM == name) { + val cls = function.containingClass + if (cls != null) { + val qualifiedName = cls.qualifiedName + return SPECIAL_FORM == qualifiedName || SPECIAL_FORM_EXT == qualifiedName + } + } + + return false + } + + private fun getTypedDictTypeForTarget(referenceTarget: PyTargetExpression, context: TypeEvalContext): PyType? { + if (isTypedDict(referenceTarget, context)) { + return PyCustomType( + TYPED_DICT, null, false, true, + getInstance(referenceTarget).dictType + ) + } + + return null + } + + private fun getMemberTypeForClassType( + context: Context, + target: PyTargetExpression?, + name: String, + resolveContext: PyResolveContext, + isInherited: Boolean, + classType: PyClassTypeImpl, + ): Ref? { + val classAttrs = + classType.resolveMember(name, target, AccessDirection.READ, resolveContext, isInherited) + if (classAttrs == null) { + return null + } + return StreamEx.of(classAttrs) + .map { obj: RatedResolveResult? -> obj!!.element } + .select(PyTargetExpression::class.java) + .filter { x: PyTargetExpression? -> + val owner = ScopeUtil.getScopeOwner(x) + owner is PyClass || owner is PyFunction + } + .map?> { x: PyTargetExpression? -> getTypeFromTypeHint(x, context) } + .collect(PyTypeUtil.toUnionFromRef()) + } + + private fun + getTypeFromTypeHint(element: T, context: Context): Ref? where T : PyAnnotationOwner?, T : PyTypeCommentOwner? { + val annotation: PyExpression? = getAnnotationValue(element!!, context.typeContext) + if (annotation != null) { + return getType(annotation, context) + } + val comment = element.typeCommentAnnotation + if (comment != null) { + return getVariableTypeCommentType(comment, element, context) + } + return null + } + + /** + * Checks that text of a comment starts with "# type:" prefix and returns trimmed type hint after it. + * The trailing part is supposed to contain type annotation in PEP 484 compatible format and an optional + * plain text comment separated from it with another "#". + * + * + * For instance, for `# type: List[int]# comment` it returns `List[int]`. + * + * + * This method cannot return an empty string. + * + * @see .getTypeCommentValueRange + */ + @JvmStatic + fun getTypeCommentValue(text: String): String? { + return PyUtilCore.getTypeCommentValue(text) + } + + /** + * Returns the corresponding text range for a type hint as returned by [.getTypeCommentValue]. + * + * @see .getTypeCommentValue + */ + fun getTypeCommentValueRange(text: String): TextRange? { + return PyUtilCore.getTypeCommentValueRange(text) + } + + private fun evaluateSuperClassesAsTypeHints(pyClass: PyClass, context: TypeEvalContext): MutableList { + val results: MutableList = ArrayList() + for (superClassExpression in getSuperClassExpressions(pyClass)) { + val type = Ref.deref(getType(superClassExpression, context)) + if (type is PyClassType) { + results.add(type) + } + } + return results + } + + private fun collectTypeParameters(cls: PyClass, context: Context): List { + if (!isGeneric(cls, context.typeContext)) { + return emptyList() + } + if (cls.typeParameterList != null) { + val typeParameters = cls.typeParameterList!!.typeParameters + return StreamEx.of(typeParameters) + .map { + getTypeParameterTypeFromTypeParameter( + it!!, + context + ) + } + .nonNull() + .toList() + } + // See https://mypy.readthedocs.io/en/stable/generics.html#defining-sub-classes-of-generic-classes + val parameterizedSuperClassExpressions = getSuperClassExpressions(cls).filterIsInstance() + val genericAsSuperClass = parameterizedSuperClassExpressions.find { + resolvesToQualifiedNames(it.operand, + context.typeContext, + GENERIC) + } + return StreamEx.of( + if (genericAsSuperClass != null) listOf(genericAsSuperClass) + else parameterizedSuperClassExpressions + ) + .map { it!!.indexExpression } + .flatMap { + val tupleExpr = it as? PyTupleExpression + if (tupleExpr != null) StreamEx.of(*tupleExpr.elements) else StreamEx.of(it) + } + .nonNull() + .map { getType(it!!, context) } + .map { Ref.deref(it) } + .flatMap { + val typeParams = PyTypeChecker.collectGenerics(it, context.typeContext) + StreamEx.of(typeParams.typeVars).append(typeParams.typeVarTuples) + .append(StreamEx.of(typeParams.paramSpecs)) + } + .select(PyTypeParameterType::class.java) + .distinct() + .toList() + } + + /** + * If the class' stub is present, return expressions in the base classes list, converting + * their saved text chunks into [PyExpressionCodeFragment] and extracting top-level expressions + * from them. Otherwise, get superclass expressions directly from AST. + */ + private fun getSuperClassExpressions(pyClass: PyClass): List { + val classStub = pyClass.stub + if (classStub == null) { + return pyClass.superClassExpressions.toList() + } + return CachedValuesManager.getCachedValue(pyClass, CachedValueProvider { + CachedValueProvider.Result.create( + classStub.superClassesText.mapNotNull { PyUtil.createExpressionFromFragment(it, pyClass) }, + PsiModificationTracker.MODIFICATION_COUNT, + ) + } + ) + } + + private fun collectTypeParametersFromTypeAliasStatement( + typeAliasStatement: PyTypeAliasStatement, + context: Context, + ): MutableList { + val typeParameterList = typeAliasStatement.typeParameterList + if (typeParameterList != null) { + val typeParameters = typeParameterList.typeParameters + return StreamEx.of(typeParameters) + .map { typeParameter: PyTypeParameter? -> + getTypeParameterTypeFromTypeParameter( + typeParameter!!, + context + ) + } + .nonNull() + .toList() + } + return mutableListOf() + } + + @JvmStatic + fun isGeneric(descendant: PyWithAncestors, context: TypeEvalContext): Boolean { + if (descendant is PyClass && descendant.typeParameterList != null || + descendant is PyClassType && descendant.pyClass.typeParameterList != null + ) { + return true + } + for (ancestor in descendant.getAncestorTypes(context)) { + if (ancestor != null) { + if (ancestor.classQName in GENERIC_CLASSES) { + return true + } + else if (ancestor is PyClassType && + ancestor.pyClass.typeParameterList != null + ) { + return true + } + } + } + return false + } + + @ApiStatus.Internal + @JvmStatic + fun getType(expression: PyExpression, context: TypeEvalContext, useFqn: Boolean): Ref? { + return staticWithCustomContext( + context, + useFqn + ) { customContext: Context? -> getType(expression, customContext!!) } + } + + @JvmStatic + fun getType(expression: PyExpression, context: TypeEvalContext): Ref? { + return staticWithCustomContext( + context + ) { customContext: Context? -> getType(expression, customContext!!) } + } + + @JvmStatic + fun getTypeForTypeHint(expression: PyExpression, context: TypeEvalContext): Ref? { + return staticWithCustomContext( + context, + true + ) { getTypeForResolvedElement(expression, null, expression, it) } + } + + private fun getType(expression: PyExpression, context: Context): Ref? { + context.getKnownType(expression)?.let { + return Ref(it) + } + for (pair in tryResolvingWithAliases(expression, context.typeContext)) { + val typeRef = getTypeForResolvedElement(expression, pair.first, pair.second!!, context) + if (typeRef != null) { + if (typeRef.get() != null) { + context.assumeType(expression, typeRef.get()!!) + } + return typeRef + } + } + return null + } + + private fun typeHasOverloadedBitwiseOr( + type: PyType, expression: PyExpression, + context: Context, + ): Boolean { + if (type !is PyClassType) { + return false + } + val typeContext = context.typeContext + val metaClassType = type.getMetaClassType(typeContext, true) + if (metaClassType == null) { + return false + } + val resolved = metaClassType + .resolveMember("__or__", expression, AccessDirection.READ, PyResolveContext.defaultContext(typeContext)) + if (resolved.isNullOrEmpty()) return false + + return StreamEx.of(resolved) + .map { it: RatedResolveResult? -> it!!.element } + .nonNull() + .noneMatch { it: PsiElement? -> getInstance(it).isBuiltin(it) } + } + + @JvmStatic + fun isBitwiseOrUnionAvailable(context: TypeEvalContext): Boolean { + val originFile = context.origin + return originFile == null || isBitwiseOrUnionAvailable(originFile) + } + + @JvmStatic + fun isBitwiseOrUnionAvailable(element: PsiElement): Boolean { + if (LanguageLevel.forElement(element).isAtLeast(LanguageLevel.PYTHON310)) return true + + val file = element.containingFile + if (file is PyFile && file.hasImportFromFuture(FutureFeature.ANNOTATIONS)) { + return file === element || PsiTreeUtil.getParentOfType( + element, + PyAnnotation::class.java, + false, + PyStatement::class.java + ) != null + } + + return false + } + + private fun getTypeForResolvedElement( + typeHint: PyExpression, + alias: PyQualifiedNameOwner?, + resolved: PsiElement, + context: Context, + ): Ref? { + if (alias != null) { + if (context.containsTypeAlias(alias)) { + // Recursive types are not yet supported + return null + } + context.addTypeAlias(alias) + } + if (resolved is PyClass && !context.addClassDeclaration(resolved)) { + // Resolving to normal classes shouldn't cause recursive evaluation of type hints, + // but constructing recursive PyTypedDictTypes will trigger that. + return null + } + try { + val typeHintFromProvider = parseTypeHint( + typeHint, + alias, + resolved, + context.typeContext + ) + if (typeHintFromProvider != null) { + return typeHintFromProvider + } + val typeFromParenthesizedExpression: Ref? = getTypeFromParenthesizedExpression(resolved, context) + if (typeFromParenthesizedExpression != null) { + return typeFromParenthesizedExpression + } + // We perform chained resolve only for actual aliases as tryResolvingWithAliases() returns the passed-in + // expression both when it's not a reference expression and when it's failed to resolve it, hence we might + // hit SOE for mere unresolved references in the latter case. + if (alias != null) { + val typeFromTypeAlias: Ref? = getTypeFromTypeAlias(alias, typeHint, resolved, context) + if (typeFromTypeAlias != null) { + return typeFromTypeAlias + } + } + val neverType: PyType? = getNeverType(resolved) + if (neverType != null) { + return Ref(neverType) + } + val unionType: Ref? = getUnionType(resolved, context) + if (unionType != null) { + return unionType + } + val intersectionType: Ref? = getIntersectionType(resolved, context) + if (intersectionType != null) { + return intersectionType + } + val concatenateType: PyType? = getConcatenateType(resolved, context) + if (concatenateType != null) { + return Ref(concatenateType) + } + val optionalType: Ref? = getOptionalType(resolved, context) + if (optionalType != null) { + return optionalType + } + val callableType: PyType? = getCallableType(resolved, context) + if (callableType != null) { + return Ref(callableType) + } + val classVarType: Ref? = unwrapTypeModifier(resolved, context, CLASS_VAR) + if (classVarType != null) { + return classVarType + } + val classObjType: Ref? = getClassObjectType(resolved, context) + if (classObjType != null) { + return classObjType + } + val finalType: Ref? = unwrapTypeModifier(resolved, context, FINAL, FINAL_EXT) + if (finalType != null) { + return finalType + } + val annotatedType: Ref? = getAnnotatedType(resolved, context) + if (annotatedType != null) { + return annotatedType + } + val requiredOrNotRequiredType: Ref? = getTypedDictSpecialItemType(resolved, context) + if (requiredOrNotRequiredType != null) { + return requiredOrNotRequiredType + } + val literalStringType: Ref? = getLiteralStringType(resolved, context) + if (literalStringType != null) { + return literalStringType + } + val literalType: Ref? = getLiteralType(resolved, context) + if (literalType != null) { + return literalType + } + val typeAliasType: Ref? = getExplicitTypeAliasType(resolved) + if (typeAliasType != null) { + return typeAliasType + } + val narrowedType: Ref? = getNarrowedType(resolved, context) + if (narrowedType != null) { + return narrowedType + } + val parameterizedType: PyType? = getParameterizedType(resolved, context) + if (parameterizedType != null) { + return Ref(parameterizedType) + } + val collection: PyType? = getCollection(resolved, context.typeContext) + if (collection != null) { + return Ref(collection) + } + val typeParameter: PyType? = getTypeParameterTypeFromDeclaration(resolved, context) + if (typeParameter != null) { + return Ref(anchorTypeParameter(typeHint, typeParameter, context)) + } + val unpackedType: PyType? = getUnpackedType(resolved, context.typeContext) + if (unpackedType != null) { + return Ref(unpackedType) + } + val typeParameterType: PyType? = getTypeParameterTypeFromTypeParameter(resolved, context) + if (typeParameterType != null) { + return Ref(typeParameterType) + } + val callableParameterListType: PyType? = getCallableParameterListType(resolved, context) + if (callableParameterListType != null) { + return Ref(callableParameterListType) + } + val stringBasedType: PyType? = getStringLiteralType(resolved, context) + if (stringBasedType != null) { + return Ref(stringBasedType) + } + val anyType: Ref? = getAnyType(resolved, context) + if (anyType != null) { + return anyType + } + val typedDictType: PyType? = getTypedDictTypeForResolvedElement( + resolved, + context.typeContext + ) + if (typedDictType != null) { + return Ref(typedDictType) + } + val selfType: Ref? = getSelfType(resolved, typeHint, context) + if (selfType != null) { + return selfType + } + val noneType: Ref? = getNoneType(typeHint, resolved) + if (noneType != null) { + return noneType + } + val newType = PyTypingNewTypeTypeProvider.getNewTypeForResolvedElement(resolved, context.typeContext) + if (newType != null) { + return Ref(newType.toInstance()) + } + val classType: Ref? = getClassType(typeHint, resolved, context) + if (classType != null) { + return classType + } + if (context.typeRepresentationMode) { + if (resolved.text == "Unknown") { + return Ref() + } + if (resolved is PyFunctionTypeRepresentation) { + val result = context.typeContext.getType(resolved) + if (result != null) { + return Ref(result) + } + } + if (resolved is PySubscriptionExpression) { + val moduleType: Ref? = getModuleType(resolved) + if (moduleType != null) { + return moduleType + } + } + } + return null + } + finally { + if (resolved is PyClass) { + context.removeClassDeclaration(resolved) + } + if (alias != null) { + context.removeTypeAlias(alias) + } + } + } + + private fun getModuleType(moduleDefinition: PySubscriptionExpression): Ref? { + val name = moduleDefinition.rootOperand.name + if (name == null || name != PyModuleTypeName) { + return null + } + val moduleReferenceExpression = moduleDefinition.indexExpression as? PyReferenceExpression ?: return null + val moduleName = moduleReferenceExpression.name ?: return null + val project = moduleDefinition.project + val moduleInitFiles = PyModuleNameIndex.findByShortName(moduleName, project, GlobalSearchScope.everythingScope(project)) + val skeletons = PythonSdkUtil.getSkeletonsRootPath(PathManager.getSystemDir().toString()) + val firstModuleInitFile = + moduleInitFiles.find { it != null && !Path.of(it.virtualFile.path).startsWith(skeletons) } + ?: moduleInitFiles.find { it != null } + ?: return null + + return Ref(PyModuleType(firstModuleInitFile)) + } + + private fun getIntersectionType(resolved: PsiElement, context: Context): Ref? { + if (resolved is PyBinaryExpression && resolved.operator === PyTokenTypes.AND) { + val left = resolved.leftExpression + val right = resolved.rightExpression + if (left == null || right == null) return null + + val leftTypeRef: Ref? = getType(left, context) + val rightTypeRef: Ref? = getType(right, context) + if (leftTypeRef == null || rightTypeRef == null) return null + + val intersection = intersection(leftTypeRef.get(), rightTypeRef.get()) + return if (intersection != null) Ref(intersection) else null + } + return null + } + + private fun getNoneType(typeHint: PyExpression, resolved: PsiElement): Ref? { + if (typeHint is PyNoneLiteralExpression || + typeHint is PyReferenceExpression && PyNames.NONE == typeHint.text + ) { + return Ref(getInstance(resolved).noneType) + } + return null + } + + private fun getCallableParameterListType(resolved: PsiElement, context: Context): PyType? { + if (resolved is PyListLiteralExpression) { + val argumentTypes = + resolved.elements.map { + Ref.deref(getType(it!!, context)) + } + return PyCallableParameterListTypeImpl( + argumentTypes.map { PyCallableParameterImpl.nonPsi(it) } + ) + } + return null + } + + private fun getNarrowedType(resolved: PsiElement, context: Context): Ref? { + if (resolved is PySubscriptionExpression) { + val names = resolveToQualifiedNames( + resolved.operand, + context.typeContext + ) + val isTypeIs = TYPE_IS in names || TYPE_IS_EXT in names + val isTypeGuard = TYPE_GUARD in names || TYPE_GUARD_EXT in names + if (isTypeIs || isTypeGuard) { + val indexTypes: MutableList = getIndexTypes(resolved, context) + if (indexTypes.size == 1) { + val narrowedType = create(resolved, isTypeIs, indexTypes[0]) + if (narrowedType != null) { + return Ref(narrowedType) + } + } + } + } + return null + } + + private fun getSelfType(resolved: PsiElement, typeHint: PyExpression, context: Context): Ref? { + if (resolved is PyQualifiedNameOwner && + (SELF == resolved.qualifiedName || + SELF_EXT == resolved.qualifiedName) + ) { + val typeHintContext: PsiElement = getStubRetainedTypeHintContext(typeHint) + + val containingClass = typeHintContext as? PyClass ?: PsiTreeUtil.getStubOrPsiParentOfType(typeHintContext, PyClass::class.java) + if (containingClass == null) return null + + val scopeClassType = PyUtil.`as`(containingClass.getType(context.typeContext), PyClassType::class.java) + if (scopeClassType == null) return null + + return Ref(PySelfType(scopeClassType).toInstance()) + } + return null + } + + private fun getTypeFromParenthesizedExpression(resolved: PsiElement, context: Context): Ref? { + if (resolved is PyParenthesizedExpression) { + val containedExpression = PyPsiUtils.flattenParens(resolved as PyExpression) + return if (containedExpression != null) getType(containedExpression, context) else null + } + return null + } + + private fun getExplicitTypeAliasType(resolved: PsiElement): Ref? { + if (resolved is PyQualifiedNameOwner) { + val qualifiedName = resolved.qualifiedName + if (TYPE_ALIAS == qualifiedName || TYPE_ALIAS_EXT == qualifiedName) { + return Ref() + } + } + return null + } + + private fun anchorTypeParameter(typeHint: PyExpression, type: PyType?, context: Context): PyType? { + val typeParamDefinitionFromStack = if (context.isTypeAliasStackEmpty) null else context.peekTypeAlias() + assert(typeParamDefinitionFromStack == null || typeParamDefinitionFromStack is PyTargetExpression) + val targetExpr = typeParamDefinitionFromStack as PyTargetExpression? + if (type is PyTypeVarTypeImpl) { + return type.withScopeOwner(getTypeParameterScope(type.name, typeHint, context)).withDeclarationElement(targetExpr) + } + if (type is PyParamSpecType) { + return type.withScopeOwner(getTypeParameterScope(type.name, typeHint, context)).withDeclarationElement(targetExpr) + } + if (type is PyTypeVarTupleTypeImpl) { + return type.withScopeOwner(getTypeParameterScope(type.name, typeHint, context)) + .withDeclarationElement(targetExpr) + } + return type + } + + private fun getClassObjectType(resolved: PsiElement, context: Context): Ref? { + if (resolved is PySubscriptionExpression) { + val operand = resolved.operand + if (resolvesToQualifiedNames( + operand, + context.typeContext, TYPE, PyNames.TYPE + ) + ) { + val indexExpr = resolved.indexExpression + if (indexExpr != null) { + if (resolvesToQualifiedNames(indexExpr, context.typeContext, ANY)) { + return Ref(getInstance(resolved).typeType) + } + return getAsClassObjectType(indexExpr, context) + } + // Map Type[Something] with unsupported type parameter to Any, instead of generic type for the class "type" + return Ref() + } + } + else if (TYPE == getQualifiedName(resolved)) { + return Ref(getInstance(resolved).typeType) + } + return null + } + + private fun getAsClassObjectType(expression: PyExpression, context: Context): Ref { + val type = Ref.deref(getType(expression, context)) + val classType = type as? PyClassType + if (classType != null && !classType.isDefinition) { + return Ref(classType.toClass()) + } + val typeVar = type as? PyTypeVarType + if (typeVar != null && !typeVar.isDefinition) { + return Ref(typeVar.toClass()) + } + val selfType = type as? PySelfType + if (selfType != null) { + return Ref(selfType.toClass()) + } + // Represent Type[Union[str, int]] internally as Union[Type[str], Type[int]] + if (type is PyUnionType && + type.members.all { it is PyClassType && !it.isDefinition } + ) { + return Ref(type.map { (it as PyClassType).toClass() }) + } + return Ref() + } + + private fun getAnyType(element: PsiElement, context: Context): Ref? { + if (ANY == getQualifiedName(element)) { + return Ref() + } + if (context.typeRepresentationMode && ANY == element.text) { + return Ref() + } + return null + } + + private fun getClassType(typeHint: PyExpression, element: PsiElement, context: Context): Ref? { + if (typeHint is PyReferenceExpression && element is PyTypedElement) { + val typeContext = context.typeContext + val type: PyType? + if (context.typeRepresentationMode && element is PyReferenceExpression) { + val qualifiedName = element.asQualifiedName() + val project = element.project + val class_ = if (qualifiedName != null) PyPsiFacade.getInstance(project) + .createClassByQName(qualifiedName.toString(), element) + else null + type = class_?.getType(typeContext) + } + else { + type = typeContext.getType(element) + } + if (type is PyClassLikeType) { + if (type.isDefinition) { + // If we're interpreting a type hint like "MyGeneric" that is not followed by a list of type arguments (e.g. MyGeneric[int]), + // we want to parameterize it with its type parameters defaults already here. + // We need this check for the type argument list because getParameterizedType() relies on getClassType() for + // getting the type corresponding to the subscription expression operand. + val stubRetainedContext: PsiElement = getStubRetainedTypeHintContext(typeHint) + if ( + type is PyClassType + && !(stubRetainedContext is PyClass || + PsiTreeUtil.getStubOrPsiParentOfType( + stubRetainedContext, + ScopeOwner::class.java + ) is PyClass) + && (typeHint.parent as? PySubscriptionExpression)?.operand != typeHint && + isGeneric(type, context.typeContext) + ) { + val parameterized = + parameterizeClassDefaultAware(type.pyClass, listOf(), context) + if (parameterized != null) { + return Ref(parameterized.toInstance()) + } + } + val instanceType: PyType = type.toInstance() + return Ref(instanceType) + } + } + } + return null + } + + private fun getOptionalType(element: PsiElement, context: Context): Ref? { + if (element is PySubscriptionExpression) { + if (resolvesToQualifiedNames(element.operand, context.typeContext, OPTIONAL)) { + val indexExpr = element.indexExpression + if (indexExpr != null) { + val typeRef: Ref? = getType(indexExpr, context) + if (typeRef != null) { + return Ref(PyUnionType.union(typeRef.get(), getInstance(element).noneType)) + } + } + return Ref() + } + } + return null + } + + private fun getLiteralStringType(resolved: PsiElement, context: Context): Ref? { + if (resolved is PyTargetExpression) { + if (resolvesToQualifiedNames( + resolved, + context.typeContext, LITERALSTRING, LITERALSTRING_EXT + ) + ) { + return Ref(create(resolved)) + } + } + + return null + } + + private fun getLiteralType(resolved: PsiElement, context: Context): Ref? { + if (resolved is PySubscriptionExpression) { + if (resolvesToQualifiedNames(resolved.operand, context, LITERAL, LITERAL_EXT)) { + return Optional + .ofNullable(resolved.indexExpression) + .map(Function { index: PyExpression? -> + PyLiteralType.fromLiteralParameter( + index!!, + context.typeContext + ) + }) + .map(Function { value: PyType? -> Ref.create(value) }) + .orElse(null) + } + } + + return null + } + + private fun getAnnotatedType(resolved: PsiElement, context: Context): Ref? { + if (resolved is PySubscriptionExpression) { + val operand = resolved.operand + if (resolvesToQualifiedNames( + operand, + context.typeContext, ANNOTATED, ANNOTATED_EXT + ) + ) { + val indexExpr = resolved.indexExpression + val type = if (indexExpr is PyTupleExpression) indexExpr.elements[0] else indexExpr + if (type != null) { + return getType(type, context) + } + } + } + + return null + } + + private fun getTypedDictSpecialItemType(resolved: PsiElement, context: Context): Ref? { + if (resolved is PySubscriptionExpression) { + val operand = resolved.operand + if (resolvesToQualifiedNames( + operand, context.typeContext, + REQUIRED, REQUIRED_EXT, + NOT_REQUIRED, NOT_REQUIRED_EXT, + READONLY, READONLY_EXT + ) + ) { + val indexExpr = resolved.indexExpression + val type = if (indexExpr is PyTupleExpression) indexExpr.elements[0] else indexExpr + if (type != null) { + return getType(type, context) + } + } + } + + return null + } + + private fun unwrapTypeModifier(resolved: PsiElement, context: Context, vararg type: String?): Ref? { + if (resolved is PySubscriptionExpression) { + if (resolvesToQualifiedNames(resolved.operand, context.typeContext, *type)) { + val indexExpr = resolved.indexExpression + if (indexExpr != null) { + return getType(indexExpr, context) + } + } + } + + return null + } + + private fun typeHintedWithName( + owner: T, + context: TypeEvalContext, + vararg names: String?, + ): Boolean where T : PyTypeCommentOwner?, T : PyAnnotationOwner? { + return names.any { it in resolveTypeHintsToQualifiedNames(owner, context) } + } + + private fun resolveTypeHintsToQualifiedNames( + owner: T, + context: TypeEvalContext, + ): Collection where T : PyTypeCommentOwner?, T : PyAnnotationOwner? { + var annotation: PyExpression? = getAnnotationValue(owner!!, context) + if (annotation is PyStringLiteralExpression) { + val annotationText = annotation.stringValue + annotation = PyUtil.createExpressionFromFragment(annotationText, owner) + if (annotation == null) return mutableListOf() + } + + if (annotation is PySubscriptionExpression) { + return resolveToQualifiedNames(annotation.operand, context) + } + else if (annotation is PyReferenceExpression) { + return resolveToQualifiedNames(annotation, context) + } + + val typeCommentValue = owner.typeCommentAnnotation + val typeComment = if (typeCommentValue == null) null else PyUtil.createExpressionFromFragment(typeCommentValue, owner) + if (typeComment is PySubscriptionExpression) { + return resolveToQualifiedNames(typeComment.operand, context) + } + else if (typeComment is PyReferenceExpression) { + return resolveToQualifiedNames(typeComment, context) + } + + return mutableListOf() + } + + fun isFinal(decoratable: PyDecoratable, context: TypeEvalContext): Boolean { + return PyKnownDecoratorUtil + .getKnownDecorators(decoratable, context) + .any { it === PyKnownDecorator.TYPING_FINAL || it === PyKnownDecorator.TYPING_FINAL_EXT } + } + + @JvmStatic + fun isFinal(owner: T, context: TypeEvalContext): Boolean where T : PyTypeCommentOwner?, T : PyAnnotationOwner? { + return PyUtil.getParameterizedCachedValue(owner!!, context) { + typeHintedWithName(owner, context, FINAL, FINAL_EXT) + } + } + + @JvmStatic + fun isClassVar(owner: T, context: TypeEvalContext): Boolean where T : PyAnnotationOwner?, T : PyTypeCommentOwner? { + return PyUtil.getParameterizedCachedValue(owner!!, context) { + typeHintedWithName(owner, context, CLASS_VAR) + } + } + + private fun resolvesToQualifiedNames(expression: PyExpression, context: Context, vararg names: String?): Boolean { + if (!context.typeRepresentationMode) return resolvesToQualifiedNames(expression, context.typeContext, *names) + if (expression !is PyReferenceExpression) return false + val qualifier = expression.qualifier ?: return false + val qName = qualifier.name + "." + expression.name + return names.any { it == qName } + } + + private fun resolvesToQualifiedNames(expression: PyExpression, context: TypeEvalContext, vararg names: String?): Boolean { + val qualifiedNames = resolveToQualifiedNames(expression, context) + return names.any { it in qualifiedNames } + } + + @ApiStatus.Internal + fun getAnnotationValue(owner: PyAnnotationOwner, context: TypeEvalContext): PyExpression? { + if (context.maySwitchToAST(owner)) { + val annotation = owner.annotation + if (annotation != null) { + return annotation.value + } + } + else { + val annotationText = owner.annotationValue + if (annotationText != null) { + return PyUtil.createExpressionFromFragment(annotationText, owner) + } + } + return null + } + + @JvmStatic + fun getStringBasedType( + contents: String, + anchor: PsiElement, + context: TypeEvalContext, + ): Ref? { + return staticWithCustomContext( + context + ) { getStringBasedType(contents, anchor, it) } + } + + private fun getStringBasedType(contents: String, anchor: PsiElement, context: Context): Ref? { + return RecursionManager.doPreventingRecursion( + anchor to contents, + true, + Computable { + val expr = PyUtil.createExpressionFromFragment(contents, anchor) + if (expr != null) getType(expr, context) else null + }) + } + + private fun getStringLiteralType(element: PsiElement, context: Context): PyType? { + if (element is PyStringLiteralExpression) { + val contents = element.stringValue + // A multiline string literal can contain a type expression unparsable without parentheses + return Ref.deref( + getStringBasedType( + if ("\n" in contents) "($contents)" else contents, + element, + context + ) + ) + } + return null + } + + private fun getVariableTypeCommentType( + contents: String, + element: PsiElement, + context: Context, + ): Ref? { + val expr = PyPsiUtils.flattenParens(PyUtil.createExpressionFromFragment(contents, element)) + if (expr != null) { + if (element is PyTargetExpression && expr is PyTupleExpression) { + // Such syntax is specific to "# type:" comments, unpacking in type hints is not allowed anywhere else + // XXX: Switches stub to AST + val topmostTarget: PyExpression? = findTopmostTarget(element) + if (topmostTarget != null) { + val targetToExpr = mapTargetsToAnnotations(topmostTarget, expr) + val typeExpr = targetToExpr[element] + if (typeExpr != null) { + return getType(typeExpr, context) + } + } + } + else { + return getType(expr, context) + } + } + return null + } + + private fun findTopmostTarget(target: PyTargetExpression): PyExpression? { + val validTargetParent = PsiTreeUtil.getParentOfType( + target, + PyForPart::class.java, + PyWithItem::class.java, + PyAssignmentStatement::class.java + ) + if (validTargetParent == null) { + return null + } + val topmostTarget = PsiTreeUtil.findPrevParent(validTargetParent, target) as? PyExpression + if (validTargetParent is PyForPart && topmostTarget !== validTargetParent.target) { + return null + } + if (validTargetParent is PyWithItem && topmostTarget !== validTargetParent.target) { + return null + } + if (validTargetParent is PyAssignmentStatement && + validTargetParent.rawTargets.indexOf(topmostTarget) < 0 + ) { + return null + } + return topmostTarget + } + + @JvmStatic + fun mapTargetsToAnnotations( + targetExpr: PyExpression, + typeExpr: PyExpression, + ): Map { + val targetsNoParen = PyPsiUtils.flattenParens(targetExpr) + val typesNoParen = PyPsiUtils.flattenParens(typeExpr) + if (targetsNoParen == null || typesNoParen == null) { + return mutableMapOf() + } + if (targetsNoParen is PySequenceExpression && typesNoParen is PySequenceExpression) { + val result = Ref>(LinkedHashMap()) + mapTargetsToExpressions(targetsNoParen, typesNoParen, result) + return if (result.isNull) mutableMapOf() + else Collections.unmodifiableMap( + result.get() + ) + } + else if (targetsNoParen is PyTargetExpression && typesNoParen !is PySequenceExpression) { + return ImmutableMap.of(targetsNoParen, typesNoParen) + } + return emptyMap() + } + + private fun mapTargetsToExpressions( + targetSequence: PySequenceExpression, + valueSequence: PySequenceExpression, + result: Ref>, + ) { + val targets = targetSequence.elements + val values = valueSequence.elements + + if (targets.size != values.size) { + result.set(null) + return + } + + for (i in targets.indices) { + val target = PyPsiUtils.flattenParens(targets[i]) + val value = PyPsiUtils.flattenParens(values[i]) + + if (target == null || value == null) { + result.set(null) + return + } + + if (target is PySequenceExpression && value is PySequenceExpression) { + mapTargetsToExpressions(target, value, result) + if (result.isNull) { + return + } + } + else if (target is PyTargetExpression && value !is PySequenceExpression) { + val map = checkNotNull(result.get()) + map[target] = value + } + else { + result.set(null) + return + } + } + } + + private fun getCallableType(resolved: PsiElement, context: Context): PyType? { + if (resolved is PySubscriptionExpression) { + if (resolvesToQualifiedNames( + resolved.operand, + context.typeContext, CALLABLE, CALLABLE_EXT + ) + ) { + val indexExpr = resolved.indexExpression + if (indexExpr is PyTupleExpression) { + val elements = indexExpr.elements + if (elements.size == 2) { + val parametersExpr = elements[0] + val returnTypeExpr = elements[1] + var returnType = Ref.deref(getType(returnTypeExpr, context)) + if (returnType is PyVariadicType) { + returnType = null + } + if (parametersExpr is PyEllipsisLiteralExpression) { + return PyCallableTypeImpl(null as PyCallableParameterVariadicType?, returnType) + } + val parametersType = Ref.deref(getType(parametersExpr, context)) + if (parametersType is PyCallableParameterListType) { + return PyCallableTypeImpl(parametersType.parameters, returnType) + } + if (parametersType is PyCallableParameterVariadicType) { + return PyCallableTypeImpl(parametersType, returnType) + } + } + } + } + } + else if (resolved is PyTargetExpression) { + if (resolvesToQualifiedNames( + resolved, + context.typeContext, CALLABLE, CALLABLE_EXT + ) + ) { + return PyCallableTypeImpl(null as PyCallableParameterListType?, null) + } + } + return null + } + + private fun getNeverType(element: PsiElement): PyType? { + val qName: String? = getQualifiedName(element) + if (qName == null) return null + if (qName in listOf(NEVER, NEVER_EXT)) { + return PyNeverType.NEVER + } + if (qName in listOf(NO_RETURN, NO_RETURN_EXT)) { + return PyNeverType.NO_RETURN + } + return null + } + + private fun getUnionType(element: PsiElement, context: Context): Ref? { + if (element is PySubscriptionExpression) { + if (resolvesToQualifiedNames(element.operand, context.typeContext, UNION)) { + val union = PyUnionType.union(getIndexTypes(element, context)) + return if (union != null) Ref(union) else null + } + } + else if (element is PyBinaryExpression && element.operator === PyTokenTypes.OR) { + val left = element.leftExpression + val right = element.rightExpression + if (left == null || right == null) return null + + val leftTypeRef: Ref? = getType(left, context) + val rightTypeRef: Ref? = getType(right, context) + if (leftTypeRef == null || rightTypeRef == null) return null + + val leftType = leftTypeRef.get() + if (leftType != null && typeHasOverloadedBitwiseOr(leftType, left, context)) return null + + val union = PyUnionType.union(leftType, rightTypeRef.get()) + return if (union != null) Ref(union) else null + } + return null + } + + private fun getConcatenateType(element: PsiElement, context: Context): PyType? { + if (element !is PySubscriptionExpression) return null + if (!resolvesToQualifiedNames(element.operand, context.typeContext, CONCATENATE, CONCATENATE_EXT)) return null + val tupleExpression = (element.indexExpression as? PyTupleExpression) ?: return null + + val arguments = tupleExpression.elements.toList() + if (arguments.size < 2) return null + + val prefixTypeExprs = arguments.subList(0, arguments.size - 1) + val prefixTypes = prefixTypeExprs.map { Ref.deref(getType(it!!, context.typeContext)) } + + val lastTypeExpr = arguments[arguments.size - 1] + val paramSpecType = if (lastTypeExpr is PyEllipsisLiteralExpression) { + null + } + else { + Ref.deref(getType(lastTypeExpr, context.typeContext)) as? PyParamSpecType + ?: return null + } + return PyConcatenateType(prefixTypes, paramSpecType) + } + + private fun getTypeParameterTypeFromDeclaration(element: PsiElement, context: Context): PyTypeParameterType? { + if (element is PyCallExpression) { + val typeParameterKind: PyAstTypeParameter.Kind? = getTypeParameterKindFromDeclaration( + element, + context.typeContext + ) + if (typeParameterKind != null) { + val arguments = element.arguments + val nameArgument = arguments.firstOrNull() as? PyStringLiteralExpression + if (nameArgument is PyStringLiteralExpression) { + val name: String = nameArgument.stringValue + val defaultExpression = element.getKeywordArgument("default") + val defaultType: Ref? = if (defaultExpression != null) getType(defaultExpression, context) else null + when (typeParameterKind) { + PyAstTypeParameter.Kind.TypeVarTuple -> { + return PyTypeVarTupleTypeImpl(name) + .withDefaultType( + (Ref.deref(defaultType) as? PyPositionalVariadicType)?.let { Ref(it) } + ) + } + + PyAstTypeParameter.Kind.TypeVar -> { + // TypeVar __init__ parameters: + // (name, *constraints, bound = None, contravariant = False, covariant = False, infer_variance = False, default = ...) + val constraints = arguments + .drop(1) + .takeWhile { it !is PyKeywordArgument } + .map { Ref.deref(getType(it, context)) } + val boundExpression = element.getKeywordArgument("bound") + val bound = if (boundExpression == null) null else Ref.deref(getType(boundExpression, context)) + val variance: PyTypeVarType.Variance = getTypeVarVarianceFromDeclaration(element) + return PyTypeVarTypeImpl(name, constraints, bound, defaultType, variance) + } + + PyAstTypeParameter.Kind.ParamSpec -> { + return PyParamSpecType(name) + .withDefaultType( + (Ref.deref(defaultType) as? PyCallableParameterVariadicType)?.let { Ref(it) } + ) + } + } + } + } + } + return null + } + + private fun getTypeVarVarianceFromDeclaration(assignedCall: PyCallExpression): PyTypeVarType.Variance { + val covariant = PyEvaluator.evaluateAsBooleanNoResolve(assignedCall.getKeywordArgument("covariant"), false) + val contravariant = PyEvaluator.evaluateAsBooleanNoResolve(assignedCall.getKeywordArgument("contravariant"), false) + val inferVariance = PyEvaluator.evaluateAsBooleanNoResolve(assignedCall.getKeywordArgument("infer_variance"), false) + + if (covariant && !contravariant) { + return PyTypeVarType.Variance.COVARIANT + } + else if (contravariant && !covariant) { + return PyTypeVarType.Variance.CONTRAVARIANT + } + else if (inferVariance) { + return PyTypeVarType.Variance.INFER_VARIANCE + } + else { + return PyTypeVarType.Variance.INVARIANT + } + } + + @ApiStatus.Internal + fun getTypeParameterKindFromDeclaration( + callExpression: PyCallExpression, + context: TypeEvalContext, + ): PyAstTypeParameter.Kind? { + val callee = callExpression.callee + if (callee != null) { + val calleeQNames = resolveToQualifiedNames(callee, context) + if (TYPE_VAR_TUPLE in calleeQNames || TYPE_VAR_TUPLE_EXT in calleeQNames) return PyAstTypeParameter.Kind.TypeVarTuple + if (TYPE_VAR in calleeQNames || TYPE_VAR_EXT in calleeQNames) return PyAstTypeParameter.Kind.TypeVar + if (PARAM_SPEC in calleeQNames || PARAM_SPEC_EXT in calleeQNames) return PyAstTypeParameter.Kind.ParamSpec + } + return null + } + + @ApiStatus.Internal + fun getTypeParameterTypeFromTypeParameter( + typeParameter: PyTypeParameter, + context: TypeEvalContext, + ): PyTypeParameterType? { + return staticWithCustomContext( + context + ) { getTypeParameterTypeFromTypeParameter(typeParameter, it) } + } + + private fun getTypeParameterTypeFromTypeParameter( + element: PsiElement, + context: Context, + ): PyTypeParameterType? { + if (element is PyTypeParameter) { + val name = element.name + if (name == null) { + return null + } + + val typeParameterOwner = ScopeUtil.getScopeOwner(element) + val scopeOwner: PyQualifiedNameOwner? = typeParameterOwner as? PyQualifiedNameOwner + + val defaultExpressionText = element.defaultExpressionText + val defaultExpression = if (defaultExpressionText != null) + PyUtil.createExpressionFromFragment(defaultExpressionText, element) + else + null + var defaultType: Ref? = null + if (defaultExpression != null) { + val defaultExprWithoutParens = PyPsiUtils.flattenParens(defaultExpression) + defaultType = if (defaultExprWithoutParens != null) + getTypePreventingRecursion(defaultExprWithoutParens, context) + else + Ref() + } + + val declarationElement = PyUtil.`as`(element, PyQualifiedNameOwner::class.java) + + when (element.kind) { + PyAstTypeParameter.Kind.TypeVar -> { + var constraints = listOf() + var boundType: PyType? = null + val boundExpressionText = element.boundExpressionText + val boundExpression = if (boundExpressionText != null) + PyPsiUtils.flattenParens(PyUtil.createExpressionFromFragment(boundExpressionText, element)) + else + null + if (boundExpression is PyTupleExpression) { + constraints = boundExpression.elements.map { + Ref.deref( + getTypePreventingRecursion(it!!, context) + ) + } + } + else if (boundExpression != null) { + boundType = Ref.deref(getTypePreventingRecursion(boundExpression, context)) + } + return PyTypeVarTypeImpl(name, constraints, boundType, defaultType, PyTypeVarType.Variance.INFER_VARIANCE) + .withScopeOwner(scopeOwner) + .withDeclarationElement(declarationElement) + } + + PyAstTypeParameter.Kind.ParamSpec -> { + return PyParamSpecType(name) + .withScopeOwner(scopeOwner) + .withDefaultType( + (Ref.deref(defaultType) as? PyCallableParameterVariadicType)?.let { Ref(it) } + ) + .withDeclarationElement(declarationElement) + } + + PyAstTypeParameter.Kind.TypeVarTuple -> { + return PyTypeVarTupleTypeImpl(name) + .withScopeOwner(scopeOwner) + .withDefaultType( + (Ref.deref(defaultType) as? PyPositionalVariadicType)?.let { Ref(it) } + ) + .withDeclarationElement(declarationElement) + } + } + } + return null + } + + private fun getTypePreventingRecursion(expression: PyExpression, context: Context): Ref? { + return RecursionManager.doPreventingRecursion?>(expression, false, Computable { getType(expression, context) }) + } + + // See https://peps.python.org/pep-0484/#scoping-rules-for-type-variables + private fun getTypeParameterScope( + name: String, + typeHint: PyExpression, + context: Context, + ): PyQualifiedNameOwner? { + if (!context.isComputeTypeParameterScopeEnabled) return null + + val typeHintContext: PsiElement = getStubRetainedTypeHintContext(typeHint) + val typeParamOwnerCandidates = + StreamEx.iterate( + typeHintContext, + { Objects.nonNull(it) }, + UnaryOperator { owner: PsiElement? -> + PsiTreeUtil.getStubOrPsiParentOfType( + owner, + ScopeOwner::class.java + ) + }) + .filter { owner: PsiElement? -> owner is PyFunction || owner is PyClass } + .select(PyQualifiedNameOwner::class.java) + .toList() + + val closestOwner = typeParamOwnerCandidates.firstOrNull() + if (closestOwner is PyFunction) { + val typeParameterType = StreamEx.of(typeParamOwnerCandidates) + .skip(1) + .map { + findSameTypeParameterInDefinition( + it, + name, + context + ) + } + .nonNull() + .findFirst() + if (typeParameterType.isPresent) { + return typeParameterType.get().scopeOwner + } + } + if (closestOwner != null) { + val prevComputeTypeParameterScope = context.setComputeTypeParameterScopeEnabled(false) + try { + return if (findSameTypeParameterInDefinition(closestOwner, name, context) != null) closestOwner else null + } + finally { + context.setComputeTypeParameterScopeEnabled(prevComputeTypeParameterScope) + } + } + + // old-style type aliases of form `ListOf = list[T]` or `ListOf: TypeAlias = list[T]` + val assignment = PsiTreeUtil.getParentOfType( + typeHintContext, + PyAssignmentStatement::class.java, + false, + PyStatement::class.java + ) + if (assignment != null) { + val assignedValue = PyPsiUtils.flattenParens(assignment.assignedValue) + if (PsiTreeUtil.isAncestor(assignedValue, typeHintContext, false)) { + val target = PyPsiUtils.flattenParens(assignment.leftHandSideExpression) + if (target is PyTargetExpression) { + val isTypeParamDeclaration = assignedValue is PyCallExpression && + getTypeParameterKindFromDeclaration(assignedValue, context.typeContext) != null + if (!isTypeParamDeclaration && PyTypingAliasStubType.looksLikeTypeHint(assignedValue!!)) { + return target + } + } + } + } + + return null + } + + private fun findSameTypeParameterInDefinition( + owner: PyQualifiedNameOwner, + name: String, + context: Context, + ): PyTypeParameterType? { + // At this moment, the definition of the TypeVar should be the type alias at the top of the stack. + // While evaluating type hints of enclosing functions' parameters, resolving to the same TypeVar + // definition shouldn't trigger the protection against recursive aliases, so we manually remove + // it from the top for the time being. + if (context.isTypeAliasStackEmpty) { + return null + } + val typeVarDeclaration = context.popTypeAlias() + assert(typeVarDeclaration is PyTargetExpression) + try { + val typeParameters: Iterable + when (owner) { + is PyClass -> { + typeParameters = collectTypeParameters(owner, context) + } + is PyFunction -> { + typeParameters = collectTypeParameters(owner, context.typeContext) + } + else -> { + typeParameters = mutableListOf() + } + } + return typeParameters.find { name == it!!.name } + } + finally { + context.pushTypeAlias(typeVarDeclaration!!) + } + } + + @ApiStatus.Internal + fun collectTypeParameters( + function: PyFunction, + context: TypeEvalContext, + ): Iterable { + return StreamEx.of(*function.parameterList.parameters) + .select(PyNamedParameter::class.java) + .map { + PyTypingTypeProvider().getParameterType( + it!!, + function, + context + ) + } + .append(PyTypingTypeProvider().getReturnType(function, context)) + .map { Ref.deref(it) } + .map { PyTypeChecker.collectGenerics(it, context) } + .flatMap { + StreamEx.of(it!!.typeVars) + .append(it.paramSpecs) + .append(it.typeVarTuples) + } + } + + private fun getStubRetainedTypeHintContext(typeHintExpression: PsiElement): PsiElement { + val containingFile = typeHintExpression.containingFile + // Values from PSI stubs and regular type comments + val fragmentOwner = containingFile.context + if (fragmentOwner != null) { + return fragmentOwner + } + else if (containingFile is PyFunctionTypeAnnotationFile || containingFile is PyTypeHintFile) { + return PyPsiUtils.getRealContext(typeHintExpression) + } + else { + return typeHintExpression + } + } + + fun getUnpackedType(element: PsiElement, context: TypeEvalContext): PyPositionalVariadicType? { + var typeRef: Ref? = getTypeFromStarExpression(element, context) + if (typeRef == null) { + typeRef = getTypeFromUnpackOperator(element, context) + } + if (typeRef == null) { + return null + } + val expressionType = typeRef.get() + if (expressionType is PyTupleType) { + return PyUnpackedTupleTypeImpl(expressionType.elementTypes, expressionType.isHomogeneous) + } + if (expressionType is PyTypeVarTupleType) { + return expressionType + } + return null + } + + private fun getTypeFromUnpackOperator(element: PsiElement, context: TypeEvalContext): Ref? { + if (element !is PySubscriptionExpression || + !resolvesToQualifiedNames(element.operand, context, UNPACK, UNPACK_EXT) + ) { + return null + } + val indexExpression = element.indexExpression + if (!(indexExpression is PyReferenceExpression || indexExpression is PySubscriptionExpression)) return null + return Ref(Ref.deref(getType(indexExpression, context))) + } + + private fun getTypeFromStarExpression(element: PsiElement, context: TypeEvalContext): Ref? { + if (element !is PyStarExpression) return null + val starredExpression = element.expression + if (!(starredExpression is PyReferenceExpression || starredExpression is PySubscriptionExpression)) return null + return Ref(Ref.deref(getType(starredExpression, context))) + } + + private fun getIndexTypes(expression: PySubscriptionExpression, context: Context): MutableList { + val types: MutableList = ArrayList() + val indexExpr = expression.indexExpression + if (indexExpr is PyTupleExpression) { + for (expr in indexExpr.elements) { + types.add(Ref.deref(getType(expr, context))) + } + } + else if (indexExpr != null) { + types.add(Ref.deref(getType(indexExpr, context))) + } + return types + } + + private fun parameterizeClassDefaultAware( + pyClass: PyClass, + actualTypeParams: List, + context: Context, + ): PyCollectionType? { + val genericDefinitionType = + RecursionManager.doPreventingRecursion(pyClass, false, Computable { + PyTypeChecker.findGenericDefinitionType( + pyClass, + context.typeContext + ) + }) + if (genericDefinitionType != null && genericDefinitionType.elementTypes.any { + it is PyTypeParameterType && it.defaultType != null + } + ) { + val parameterizedType = PyTypeChecker.parameterizeType(genericDefinitionType, actualTypeParams, context.typeContext) + if (parameterizedType is PyCollectionType) { + return parameterizedType + } + } + return null + } + + private fun getTypeFromTypeAlias( + alias: PyQualifiedNameOwner, + typeHint: PsiElement, + element: PsiElement, + context: Context, + ): Ref? { + if (element is PyExpression) { + if (alias is PyTypeAliasStatement) { + return getTypeFromTypeAliasStatement(alias, typeHint, element, context) + } + + val assignedTypeRef: Ref? = getType(element, context) + if (assignedTypeRef != null) { + val assignedType = assignedTypeRef.get() + if (assignedType == null) { + return assignedTypeRef + } + if (typeHint is PySubscriptionExpression) { + val indexTypes: MutableList = getIndexTypes(typeHint, context) + return Ref(PyTypeChecker.parameterizeType(assignedType, indexTypes, context.typeContext)) + } + if (typeHint is PyReferenceExpression) { + if (assignedType !is PyTypeParameterType) { + val typeAliasTypeParams = + PyTypeChecker.collectGenerics(assignedType, context.typeContext).allTypeParameters + if (!typeAliasTypeParams.isEmpty()) { + return Ref( + PyTypeChecker.parameterizeType( + assignedType, + mutableListOf(), + context.typeContext + ) + ) + } + return Ref(assignedType) + } + } + } + } + return null + } + + private fun getTypeFromTypeAliasStatement( + typeAliasStatement: PyTypeAliasStatement, + typeHint: PsiElement, + assignedExpression: PyExpression, + context: Context, + ): Ref? { + val assignedTypeRef: Ref? = getType(assignedExpression, context) + if (assignedTypeRef != null) { + val assignedType = assignedTypeRef.get() + if (assignedType == null) { + return assignedTypeRef + } + val indexTypes = if (typeHint is PySubscriptionExpression) + getIndexTypes(typeHint, context) + else mutableListOf() + + val typeAliasTypeParams: MutableList = + collectTypeParametersFromTypeAliasStatement(typeAliasStatement, context) + if (!typeAliasTypeParams.isEmpty()) { + val substitutions = + PyTypeChecker.mapTypeParametersToSubstitutions( + typeAliasTypeParams, + indexTypes, + PyTypeParameterMapping.Option.USE_DEFAULTS, + PyTypeParameterMapping.Option.MAP_UNMATCHED_EXPECTED_TYPES_TO_ANY + ) + + return if (substitutions != null) Ref( + PyTypeChecker.substitute( + assignedType, + substitutions, + context.typeContext + ) + ) + else null + } + return assignedTypeRef + } + return null + } + + private fun getParameterizedType(element: PsiElement, context: Context): PyType? { + if (element is PySubscriptionExpression) { + val operand = element.operand + val indexExpr = element.indexExpression + if (indexExpr != null) { + val operandType = Ref.deref(getType(operand, context)) + val indexTypes: MutableList = getIndexTypes(element, context) + if (operandType != null) { + if (operandType is PyClassType) { + if (operandType !is PyTupleType && PyNames.TUPLE == operandType.pyClass.qualifiedName) { + if (indexExpr is PyTupleExpression) { + val elements = indexExpr.elements + if (elements.size == 2 && elements[1] is PyEllipsisLiteralExpression) { + return PyTupleType.createHomogeneous(element, indexTypes[0]) + } + } + return PyTupleType.create(element, indexTypes) + } + + if (isGeneric(operandType, context.typeContext)) { + parameterizeClassDefaultAware(operandType.pyClass, indexTypes, context)?.let { + return it.toInstance() + } + } + else { + return null + } + } + return PyTypeChecker.parameterizeType(operandType, indexTypes, context.typeContext) + } + } + } + return null + } + + private fun getCollection(element: PsiElement, context: TypeEvalContext): PyType? { + val typingName: String? = getQualifiedName(element) + + val builtinName: String? = BUILTIN_COLLECTION_CLASSES[typingName] + if (builtinName != null) return PyTypeParser.getTypeByName(element, builtinName, context) + + val collectionName: String? = COLLECTIONS_CLASSES[typingName] + if (collectionName != null) return PyTypeParser.getTypeByName(element, collectionName, context) + + return null + } + + private fun tryResolving(expression: PyExpression, context: TypeEvalContext): List { + return tryResolvingWithAliases(expression, context).map { it.second!! } + } + + private fun tryResolvingWithAliases( + expression: PyExpression, + context: TypeEvalContext, + ): List> { + val elements: MutableList> = ArrayList() + if (expression is PyReferenceExpression) { + val results: MutableList + if (context.maySwitchToAST(expression)) { + val resolveContext = PyResolveContext.defaultContext(context) + results = PyUtil.multiResolveTopPriority(expression, resolveContext) + } + else { + results = tryResolvingOnStubs(expression, context) + } + for (element in results) { + val cls = PyUtil.turnConstructorIntoClass(element as? PyFunction) + if (cls != null) { + elements.add(null to cls) + continue + } + val name: String? = if (element != null) getQualifiedName(element) else null + if (name != null && name in OPAQUE_NAMES) { + elements.add(null to element) + continue + } + // Presumably, a TypeVar definition or a type alias + if (element is PyTargetExpression) { + val assignedValue = PyTypingAliasStubType.getAssignedValueStubLike(element) + if (assignedValue != null) { + elements.add(element to assignedValue) + continue + } + } + if (element is PyTypeAliasStatement) { + val assignedValue: PyExpression? + if (context.maySwitchToAST(element)) { + assignedValue = element.typeExpression + } + else { + val assignedTypeText = element.typeExpressionText + assignedValue = + if (assignedTypeText != null) PyUtil.createExpressionFromFragment(assignedTypeText, element) else null + } + if (assignedValue != null) { + elements.add(element to assignedValue) + continue + } + } + if (element != null) { + elements.add(null to element) + } + } + } + if (expression is PySubscriptionExpression) { + // Possibly a parameterized type alias + val operandExpression = expression.operand + val results = tryResolvingWithAliases(operandExpression, context) + for (pair in results) { + // If the parameterized type is a type alias + if (pair.first != null && pair.second != null) { + elements.add(pair.first to pair.second) + } + } + } + return elements.ifEmpty { listOf(null to expression) } + } + + private fun tryResolvingOnStubs( + expression: PyReferenceExpression, + context: TypeEvalContext, + ): MutableList { + val qualifiedName = expression.asQualifiedName() + val pyFile = PyUtil.`as`(FileContextUtil.getContextFile(expression), PyFile::class.java) + + val anchor = expression.containingFile.context + val scopeOwner: ScopeOwner? + + when (anchor) { + null -> { + scopeOwner = pyFile + } + is ScopeOwner -> { + scopeOwner = anchor + } + else -> { + scopeOwner = ScopeUtil.getScopeOwner(anchor) + } + } + + if (scopeOwner != null && qualifiedName != null) { + return PyResolveUtil.resolveQualifiedNameInScope(qualifiedName, scopeOwner, context) + } + return mutableListOf(expression) + } + + fun resolveToQualifiedNames(expression: PyExpression, context: TypeEvalContext): Collection { + return buildSet { + for (resolved in tryResolving(expression, context)) { + val name: String? = getQualifiedName(resolved) + if (name != null) { + add(name) + } + } + } + } + + private fun getQualifiedName(element: PsiElement): String? { + if (element is PyQualifiedNameOwner) { + return element.qualifiedName + } + return null + } + + @JvmStatic + fun toAsyncIfNeeded(function: PyFunction, returnType: PyType?): PyType? { + if (function.isAsync && function.isAsyncAllowed) { + if (!function.isGenerator) { + return wrapInCoroutineType(returnType, function) + } + val desc = GeneratorTypeDescriptor.fromGenerator(returnType) + if (desc != null) { + val classType = PyPsiFacade.getInstance(function.project).createClassByQName(ASYNC_GENERATOR, function) + val generics = listOf(desc.yieldType, desc.sendType) + return if (classType != null) PyCollectionTypeImpl(classType, false, generics) else null + } + } + return returnType + } + + /** + * Bound narrowed types shouldn't leak out of its scope, since it is bound to a particular call site. + */ + @JvmStatic + fun removeNarrowedTypeIfNeeded(type: PyType?): PyType? { + if (type is PyNarrowedType && type.isBound()) { + return getInstance(type.original).boolType + } + else { + return type + } + } + + private fun wrapInCoroutineType(returnType: PyType?, resolveAnchor: PsiElement): PyType? { + val coroutine = PyPsiFacade.getInstance(resolveAnchor.project).createClassByQName(COROUTINE, resolveAnchor) + return if (coroutine != null) PyCollectionTypeImpl(coroutine, false, listOf(null, null, returnType)) else null + } + + @JvmStatic + fun wrapInGeneratorType( + elementType: PyType?, + sendType: PyType?, + returnType: PyType?, + anchor: PsiElement, + ): PyType? { + val generator = PyPsiFacade.getInstance(anchor.project).createClassByQName(GENERATOR, anchor) + return if (generator != null) PyCollectionTypeImpl( + generator, + false, + listOf(elementType, sendType, returnType) + ) + else null + } + + @JvmStatic + fun unwrapCoroutineReturnType(coroutineType: PyType?): Ref? { + val genericType = PyUtil.`as`(coroutineType, PyCollectionType::class.java) + + if (genericType != null) { + val qName = genericType.classQName + + if (AWAITABLE == qName) { + return Ref(ContainerUtil.getOrElse(genericType.elementTypes, 0, null)) + } + + if (COROUTINE == qName) { + return Ref(ContainerUtil.getOrElse(genericType.elementTypes, 2, null)) + } + } + + return null + } + + @JvmStatic + fun coroutineOrGeneratorElementType(coroutineOrGeneratorType: PyType?): Ref? { + val genericType = PyUtil.`as`(coroutineOrGeneratorType, PyCollectionType::class.java) + val classType = PyUtil.`as`(coroutineOrGeneratorType, PyClassType::class.java) + + if (genericType != null && classType != null) { + val qName = classType.classQName + + if (AWAITABLE == qName) { + return Ref(ContainerUtil.getOrElse(genericType.elementTypes, 0, null)) + } + + if (ArrayUtil.contains(qName, COROUTINE, GENERATOR)) { + return Ref(ContainerUtil.getOrElse(genericType.elementTypes, 2, null)) + } + } + + return null + } + + /** + * Checks whether the given assignment is type hinted with `typing.TypeAlias`. + * + * + * It can be done either with a variable annotation or a type comment. + */ + @JvmStatic + fun isExplicitTypeAlias(assignment: PyAssignmentStatement, context: TypeEvalContext): Boolean { + val target = assignment.targets.firstOrNull() as? PyTargetExpression ?: return false + return isExplicitTypeAlias(target, context) + } + + fun isExplicitTypeAlias(targetExpression: PyTargetExpression, context: TypeEvalContext): Boolean { + val annotationValue: PyExpression? = getAnnotationValue(targetExpression, context) + if (annotationValue is PyReferenceExpression) { + return resolvesToQualifiedNames(annotationValue, context, TYPE_ALIAS, TYPE_ALIAS_EXT) + } + val typeCommentAnnotation = targetExpression.typeCommentAnnotation + if (typeCommentAnnotation != null) { + val commentValue = PyUtil.createExpressionFromFragment(typeCommentAnnotation, targetExpression) + if (commentValue is PyReferenceExpression) { + return resolvesToQualifiedNames(commentValue, context, TYPE_ALIAS, TYPE_ALIAS_EXT) + } + } + return false + } + + /** + * Detects whether the given element belongs to a self-evident type hint. Namely, these are: + * + * * function and variable annotations + * * type comments + * * explicit type aliases marked with `TypeAlias` + * + * Note that `element` can belong to their AST directly or be a part of an injection inside one of such elements. + */ + @JvmStatic + fun isInsideTypeHint(element: PsiElement, context: TypeEvalContext): Boolean { + val realContext = PyPsiUtils.getRealContext(element) + + if (PsiTreeUtil.getParentOfType(realContext, PyAnnotation::class.java, false, PyStatement::class.java) != null) { + return true + } + + val comment = PsiTreeUtil.getParentOfType(realContext, PsiComment::class.java, false, PyStatement::class.java) + if (comment != null && getTypeCommentValue(comment.text) != null) { + return true + } + + val assignment = PsiTreeUtil.getParentOfType( + realContext, + PyAssignmentStatement::class.java, + false, + PyStatement::class.java + ) + if (assignment != null && + PsiTreeUtil.isAncestor(assignment.assignedValue, realContext, false) && + isExplicitTypeAlias(assignment, context) + ) { + return true + } + + val typeAlias = PsiTreeUtil.getParentOfType( + realContext, + PyTypeAliasStatement::class.java, + false, + PyStatement::class.java + ) + return typeAlias != null && PsiTreeUtil.isAncestor(typeAlias.typeExpression, realContext, false) + } + + private fun staticWithCustomContext(context: TypeEvalContext, delegate: (Context) -> T): T { + return staticWithCustomContext(context, false, delegate) + } + + private fun staticWithCustomContext( + context: TypeEvalContext, + useFqn: Boolean, + delegate: (Context) -> T, + ): T { + var customContext = context.processingContext.get(TYPE_HINT_EVAL_CONTEXT) + val firstEntrance = customContext == null + if (firstEntrance) { + customContext = Context(context, useFqn) + context.processingContext.put(TYPE_HINT_EVAL_CONTEXT, customContext) + } + try { + return delegate(customContext) + } + finally { + if (firstEntrance) { + context.processingContext.put(TYPE_HINT_EVAL_CONTEXT, null) + } + } } } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/inspections/PyFinalInspection.kt b/python/python-psi-impl/src/com/jetbrains/python/inspections/PyFinalInspection.kt index a6642a6342b3..728f4427a46e 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/inspections/PyFinalInspection.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/inspections/PyFinalInspection.kt @@ -18,14 +18,7 @@ import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyFunctionTypeA import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyParameterTypeList import com.jetbrains.python.codeInsight.parseDataclassParameters import com.jetbrains.python.codeInsight.typeHints.PyTypeHintFile -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.CLASS_VAR -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.FINAL -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.FINAL_EXT -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.getFunctionTypeAnnotation -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.getReturnTypeAnnotation -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.isFinal -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.isInsideTypeHint -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.resolveToQualifiedNames +import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider import com.jetbrains.python.psi.PyAnnotation import com.jetbrains.python.psi.PyAnnotationOwner import com.jetbrains.python.psi.PyAugAssignmentStatement @@ -139,14 +132,14 @@ class PyFinalInspection : PyInspection() { registerProblem(node.nameIdentifier, PyPsiBundle.message("INSP.final.non.method.function.could.not.be.marked.as.final")) } - getFunctionTypeAnnotation(node)?.let { comment -> + PyTypingTypeProvider.getFunctionTypeAnnotation(node)?.let { comment -> if (comment.parameterTypeList.parameterTypes.any { resolvesToFinal(if (it is PySubscriptionExpression) it.operand else it) }) { registerProblem(node.typeComment, PyPsiBundle.message("INSP.final.final.could.not.be.used.in.annotations.for.function.parameters")) } } - getReturnTypeAnnotation(node, myTypeEvalContext)?.let { + PyTypingTypeProvider.getReturnTypeAnnotation(node, myTypeEvalContext)?.let { if (resolvesToFinal(if (it is PySubscriptionExpression) it.operand else it)) { registerProblem(node.typeComment ?: node.annotation, PyPsiBundle.message("INSP.final.final.could.not.be.used.in.annotation.for.function.return.value")) @@ -463,7 +456,7 @@ class PyFinalInspection : PyInspection() { return } - if (isInsideTypeHint(node, myTypeEvalContext) && resolvesToFinal(node)) { + if (PyTypingTypeProvider.isInsideTypeHint(node, myTypeEvalContext) && resolvesToFinal(node)) { registerProblem(node, PyPsiBundle.message("INSP.final.final.could.only.be.used.as.outermost.type")) } } @@ -486,20 +479,21 @@ class PyFinalInspection : PyInspection() { ) } - private fun isFinal(decoratable: PyDecoratable) = isFinal(decoratable, myTypeEvalContext) + private fun isFinal(decoratable: PyDecoratable) = PyTypingTypeProvider.isFinal(decoratable, myTypeEvalContext) private fun isFinal(node: T): Boolean where T : PyAnnotationOwner, T : PyTypeCommentOwner { - return isFinal(node, myTypeEvalContext) + return PyTypingTypeProvider.isFinal(node, myTypeEvalContext) } private fun resolvesToFinal(expression: PyExpression?): Boolean { return expression is PyReferenceExpression && - resolveToQualifiedNames(expression, myTypeEvalContext).any { it == FINAL || it == FINAL_EXT } + PyTypingTypeProvider.resolveToQualifiedNames(expression, myTypeEvalContext) + .any { it == PyTypingTypeProvider.FINAL || it == PyTypingTypeProvider.FINAL_EXT } } private fun resolvesToClassVar(expression: PyExpression): Boolean { return (expression is PyReferenceExpression) && - resolveToQualifiedNames(expression, myTypeEvalContext).any { it == CLASS_VAR } + PyTypingTypeProvider.resolveToQualifiedNames(expression, myTypeEvalContext).any { it == PyTypingTypeProvider.CLASS_VAR } } private fun resolvesToClassVarFinal(expression: PyExpression?): Boolean { diff --git a/python/python-psi-impl/src/com/jetbrains/python/inspections/PyTypeCheckerInspection.java b/python/python-psi-impl/src/com/jetbrains/python/inspections/PyTypeCheckerInspection.java index 9476ffadb089..40bf84a87f07 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/inspections/PyTypeCheckerInspection.java +++ b/python/python-psi-impl/src/com/jetbrains/python/inspections/PyTypeCheckerInspection.java @@ -205,7 +205,7 @@ public class PyTypeCheckerInspection extends PyInspection { final var annotatedGeneratorDesc = getGeneratorDescriptorFromAnnotation(function, node); if (annotatedGeneratorDesc == null) return; - checkYieldType(annotatedGeneratorDesc.yieldType(), node, function); + checkYieldType(annotatedGeneratorDesc.yieldType, node, function); } private void visitDelegatingYieldExpression(@NotNull PyYieldExpression node, @NotNull PyFunction function) { @@ -218,7 +218,7 @@ public class PyTypeCheckerInspection extends PyInspection { if (delegateType == null) return; var delegateDesc = GeneratorTypeDescriptor.fromGeneratorOrProtocol(delegateType, myTypeEvalContext); - if (delegateDesc != null && delegateDesc.isAsync()) { + if (delegateDesc != null && delegateDesc.isAsync) { String delegateName = PythonDocumentationProvider.getTypeName(delegateType, myTypeEvalContext); registerProblem(yieldExpr, PyPsiBundle.message("INSP.type.checker.yield.from.async.generator", delegateName)); return; @@ -229,13 +229,13 @@ public class PyTypeCheckerInspection extends PyInspection { final var annotatedGeneratorDesc = getGeneratorDescriptorFromAnnotation(function, node); if (annotatedGeneratorDesc == null) return; - if (checkYieldType(annotatedGeneratorDesc.yieldType(), node, function)) return; + if (checkYieldType(annotatedGeneratorDesc.yieldType, node, function)) return; // Reversed because SendType is contravariant - final PyType expectedSendType = annotatedGeneratorDesc.sendType(); - if (delegateDesc != null && !PyTypeChecker.match(delegateDesc.sendType(), expectedSendType, myTypeEvalContext)) { + final PyType expectedSendType = annotatedGeneratorDesc.sendType; + if (delegateDesc != null && !PyTypeChecker.match(delegateDesc.sendType, expectedSendType, myTypeEvalContext)) { String expectedName = PythonDocumentationProvider.getVerboseTypeName(expectedSendType, myTypeEvalContext); - String actualName = PythonDocumentationProvider.getTypeName(delegateDesc.sendType(), myTypeEvalContext); + String actualName = PythonDocumentationProvider.getTypeName(delegateDesc.sendType, myTypeEvalContext); registerProblem(yieldExpr, PyPsiBundle.message("INSP.type.checker.yield.from.send.type.mismatch", expectedName, actualName)); } } @@ -286,7 +286,7 @@ public class PyTypeCheckerInspection extends PyInspection { if (function.isGenerator()) { final var generatorDesc = GeneratorTypeDescriptor.fromGeneratorOrProtocol(returnType, typeEvalContext); if (generatorDesc != null) { - return generatorDesc.returnType(); + return generatorDesc.returnType; } return null; } @@ -439,7 +439,7 @@ public class PyTypeCheckerInspection extends PyInspection { if (node.isGenerator()) { final var generatorDesc = GeneratorTypeDescriptor.fromGeneratorOrProtocol(annotatedType, myTypeEvalContext); final boolean shouldBeAsync = node.isAsync() && node.isAsyncAllowed(); - final boolean wrongSyncAsync = generatorDesc != null && generatorDesc.isAsync() != shouldBeAsync; + final boolean wrongSyncAsync = generatorDesc != null && generatorDesc.isAsync != shouldBeAsync; final PyType inferredType = node.getInferredReturnType(myTypeEvalContext); if (wrongSyncAsync || (generatorDesc == null && !PyTypeChecker.match(annotatedType, inferredType, myTypeEvalContext))) { diff --git a/python/python-psi-impl/src/com/jetbrains/python/inspections/PyTypeHintsInspection.kt b/python/python-psi-impl/src/com/jetbrains/python/inspections/PyTypeHintsInspection.kt index 178b88120e24..40d8780b84f6 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/inspections/PyTypeHintsInspection.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/inspections/PyTypeHintsInspection.kt @@ -36,7 +36,6 @@ import com.jetbrains.python.codeInsight.imports.AddImportHelper import com.jetbrains.python.codeInsight.imports.AddImportHelper.ImportPriority import com.jetbrains.python.codeInsight.typeHints.PyTypeHintFile import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.isBitwiseOrUnionAvailable import com.jetbrains.python.documentation.PythonDocumentationProvider import com.jetbrains.python.inspections.quickfix.PyUnpackTypeVarTupleQuickFix import com.jetbrains.python.psi.FutureFeature @@ -709,7 +708,7 @@ class PyTypeHintsInspection : PyInspection() { private fun checkInstanceAndClassChecksOn(base: PyExpression) { if (base is PyBinaryExpression && base.operator == PyTokenTypes.OR) { - if (isBitwiseOrUnionAvailable(base)) { + if (PyTypingTypeProvider.isBitwiseOrUnionAvailable(base)) { val left = base.leftExpression val right = base.rightExpression if (left != null) checkInstanceAndClassChecksOn(left) @@ -820,7 +819,7 @@ class PyTypeHintsInspection : PyInspection() { PyTypingTypeProvider.UNION, PyTypingTypeProvider.OPTIONAL, -> { - if (!isBitwiseOrUnionAvailable(base)) { + if (!PyTypingTypeProvider.isBitwiseOrUnionAvailable(base)) { registerParametrizedGenericsProblem(qName, base) } else if (base is PySubscriptionExpression) { diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyYieldExpressionImpl.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyYieldExpressionImpl.java index 8484c9864e33..e923a2314a78 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyYieldExpressionImpl.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyYieldExpressionImpl.java @@ -31,7 +31,7 @@ public class PyYieldExpressionImpl extends PyElementImpl implements PyYieldExpre final PyType type = e != null ? context.getType(e) : null; var generatorDesc = PyTypingTypeProvider.GeneratorTypeDescriptor.fromGeneratorOrProtocol(type, context); if (generatorDesc != null) { - return generatorDesc.returnType(); + return generatorDesc.returnType; } return PyBuiltinCache.getInstance(this).getNoneType(); } @@ -58,7 +58,7 @@ public class PyYieldExpressionImpl extends PyElementImpl implements PyYieldExpre var returnType = context.getReturnType(function); var generatorDesc = PyTypingTypeProvider.GeneratorTypeDescriptor.fromGeneratorOrProtocol(returnType, context); if (generatorDesc != null) { - return generatorDesc.sendType(); + return generatorDesc.sendType; } } } @@ -68,7 +68,7 @@ public class PyYieldExpressionImpl extends PyElementImpl implements PyYieldExpre final PyType type = e != null ? context.getType(e) : null; var generatorDesc = PyTypingTypeProvider.GeneratorTypeDescriptor.fromGeneratorOrProtocol(type, context); if (generatorDesc != null) { - return generatorDesc.sendType(); + return generatorDesc.sendType; } return PyBuiltinCache.getInstance(this).getNoneType(); } diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyExpectedTypeJudgement.kt b/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyExpectedTypeJudgement.kt index 441e237c7cfc..f10c7f094f03 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyExpectedTypeJudgement.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyExpectedTypeJudgement.kt @@ -439,7 +439,7 @@ object PyExpectedTypeJudgement { val returnType = ctx.getReturnType(funScope) val generatorDescriptor = PyTypingTypeProvider.GeneratorTypeDescriptor.fromGenerator(returnType) - val yieldType = generatorDescriptor?.yieldType() + val yieldType = generatorDescriptor?.yieldType if (parent.isDelegating) { return createIterableType(expr, yieldType) } @@ -455,7 +455,7 @@ object PyExpectedTypeJudgement { if (funScope.isAsync) { return PyTypingTypeProvider.unwrapCoroutineReturnType(returnType)?.get() } - val generatorReturnType = PyTypingTypeProvider.GeneratorTypeDescriptor.fromGenerator(returnType)?.returnType() + val generatorReturnType = PyTypingTypeProvider.GeneratorTypeDescriptor.fromGenerator(returnType)?.returnType return generatorReturnType ?: returnType } diff --git a/python/src/com/jetbrains/python/inlayHints/PyTypeInlayHintsProvider.kt b/python/src/com/jetbrains/python/inlayHints/PyTypeInlayHintsProvider.kt index cce205bc351b..c6dc646b5b56 100644 --- a/python/src/com/jetbrains/python/inlayHints/PyTypeInlayHintsProvider.kt +++ b/python/src/com/jetbrains/python/inlayHints/PyTypeInlayHintsProvider.kt @@ -14,8 +14,7 @@ import com.intellij.codeInsight.hints.declarative.impl.PresentationTreeBuilderIm import com.intellij.openapi.editor.Editor import com.intellij.psi.PsiElement import com.intellij.psi.PsiFile -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.REVEAL_TYPE -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.REVEAL_TYPE_EXT +import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider import com.jetbrains.python.documentation.PythonDocumentationProvider import com.jetbrains.python.psi.PyCallExpression import com.jetbrains.python.psi.PyFunction @@ -55,7 +54,8 @@ class PyTypeInlayHintsProvider : InlayHintsProvider { val callable = element.multiResolveCalleeFunction(resolveContext).singleOrNull() val typeEvalContext = resolveContext.typeEvalContext - if (callable is PyFunction && callable.qualifiedName in listOf(REVEAL_TYPE, REVEAL_TYPE_EXT)) { + if (callable is PyFunction && callable.qualifiedName in listOf(PyTypingTypeProvider.REVEAL_TYPE, + PyTypingTypeProvider.REVEAL_TYPE_EXT)) { val args = element.getArguments() if (args.size != 1) return diff --git a/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt b/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt index f59b15d02d69..9bedada9e77e 100644 --- a/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt +++ b/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestFixtureReferenceContributor.kt @@ -20,8 +20,7 @@ import com.intellij.util.ProcessingContext import com.intellij.util.containers.ContainerUtil import com.intellij.util.containers.toArray import com.jetbrains.python.BaseReference -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.COROUTINE -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.GENERATOR +import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider import com.jetbrains.python.psi.PyArgumentList import com.jetbrains.python.psi.PyCallExpression import com.jetbrains.python.psi.PyDecorator @@ -143,10 +142,10 @@ class PyTextFixtureTypeProvider : PyTypeProviderBase() { val classType = PyUtil.`as`(type, PyClassType::class.java) if (genericType != null && classType != null) { val qName = classType.getClassQName() - if (ArrayUtil.contains(qName, "typing.Awaitable", GENERATOR)) { + if (ArrayUtil.contains(qName, "typing.Awaitable", PyTypingTypeProvider.GENERATOR)) { return Ref.create(ContainerUtil.getOrElse(genericType.getElementTypes(), 0, null)) } - if (COROUTINE == qName) { + if (PyTypingTypeProvider.COROUTINE == qName) { return Ref.create(ContainerUtil.getOrElse(genericType.getElementTypes(), 2, null)) } }