diff --git a/python/python-psi-api/resources/intellij.python.psi.xml b/python/python-psi-api/resources/intellij.python.psi.xml index 44c9a96a799a..46fa08063dce 100644 --- a/python/python-psi-api/resources/intellij.python.psi.xml +++ b/python/python-psi-api/resources/intellij.python.psi.xml @@ -41,5 +41,8 @@ + diff --git a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java index 0f95a6641912..0e08b2f16fb1 100644 --- a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java +++ b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java @@ -1,6 +1,8 @@ // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. package com.jetbrains.python.psi.types; +import com.intellij.openapi.diagnostic.Logger; +import com.intellij.openapi.progress.ProcessCanceledException; import com.intellij.openapi.project.Project; import com.intellij.openapi.roots.ProjectFileIndex; import com.intellij.openapi.util.Pair; @@ -51,9 +53,12 @@ public sealed class TypeEvalContext { private final ThreadLocal myProcessingContext = ThreadLocal.withInitial(ProcessingContext::new); protected final Map myEvaluated = CollectionFactory.createConcurrentSoftValueMap(); + public final Map myExternalEvaluated = CollectionFactory.createConcurrentSoftValueMap(); protected final Map myEvaluatedReturn = CollectionFactory.createConcurrentSoftValueMap(); protected final Map, PyType> contextTypeCache = CollectionFactory.createConcurrentSoftValueMap(); + protected static final Logger logger = Logger.getInstance(TypeEvalContext.class); + private TypeEvalContext(boolean allowDataFlow, boolean allowStubToAST, boolean allowCallContext, @Nullable PsiFile origin) { this(new TypeEvalConstraints(allowDataFlow, allowStubToAST, allowCallContext, origin)); } @@ -216,6 +221,11 @@ public sealed class TypeEvalContext { return this instanceof AssumptionContext; } + @ApiStatus.Internal + public boolean isKnown(PyTypedElement element) { + return getKnownType(element) != null; + } + protected @Nullable PyType getKnownType(final @NotNull PyTypedElement element) { if (element instanceof PyInstantTypeProvider) { return element.getType(this, Key.INSTANCE); @@ -225,6 +235,11 @@ public sealed class TypeEvalContext { assertValid(cachedType, element); return cachedType; } + final PyType cachedExternalType = myExternalEvaluated.get(element); + if (cachedExternalType != null) { + assertValid(cachedExternalType, element); + return cachedExternalType; + } return null; } @@ -237,6 +252,11 @@ public sealed class TypeEvalContext { return null; } + @ApiStatus.Experimental + public void putExternalType(PyTypedElement element, PyType type) { + myExternalEvaluated.put(element, type == null ? PyNullType.INSTANCE : type); + } + private static boolean isLibraryElement(@NotNull PsiElement element) { PsiFile containingFile = element.getContainingFile(); VirtualFile vFile = containingFile == null ? null : containingFile.getOriginalFile().getVirtualFile(); @@ -273,6 +293,24 @@ public sealed class TypeEvalContext { Pair.create(element, this), false, () -> { + // Try external providers first + for (var provider : TypeEvalExternalTypeProvider.EP_NAME.getExtensionList()) { + try { + var provided = provider.provideType(element, this); + if (provided != null) { + var type = provided.get(); + myExternalEvaluated.put(element, type == null ? PyNullType.INSTANCE : type); + return type; + } + } + catch (ProcessCanceledException e) { + throw e; + } + catch (Exception e) { + logger.warn("Exception during external type provider " + provider.getClass().getName(), e); + } + } + PyType type = element.getType(this, Key.INSTANCE); assertValid(type, element); myEvaluated.put(element, type == null ? PyNullType.INSTANCE : type); diff --git a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalExternalTypeProvider.kt b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalExternalTypeProvider.kt new file mode 100644 index 000000000000..b46fc04f625c --- /dev/null +++ b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalExternalTypeProvider.kt @@ -0,0 +1,21 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.psi.types + +import com.intellij.openapi.extensions.ExtensionPointName +import com.intellij.openapi.util.Ref +import com.intellij.psi.PsiElement +import com.jetbrains.python.psi.PyTypedElement + +/** + * An external type provider invoked by [TypeEvalContext] to obtain a type from a separate engine. + * Implementations may return `null` if they cannot provide a type for the given element. + */ +interface TypeEvalExternalTypeProvider { + fun provideType(element: PyTypedElement, context: TypeEvalContext): Ref? + + companion object { + @JvmField + val EP_NAME: ExtensionPointName = + ExtensionPointName.create("Pythonid.typeEvalExternalTypeProvider") + } +} diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java index 9166c7b304f0..8a6ae5d4165e 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java @@ -783,6 +783,11 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext< 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 getType(@NotNull PyExpression expression, @NotNull TypeEvalContext context) { return staticWithCustomContext(context, customContext -> getType(expression, customContext)); } @@ -978,6 +983,9 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext< if (classType != null) { return classType; } + if (context.myUseFqn && resolved.getText().equals("Unknown")) { + return Ref.create(); + } return null; } finally { @@ -1176,7 +1184,7 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext< private static @Nullable Ref getLiteralType(@NotNull PsiElement resolved, @NotNull Context context) { if (resolved instanceof PySubscriptionExpression subscriptionExpr) { - if (resolvesToQualifiedNames(subscriptionExpr.getOperand(), context.getTypeContext(), LITERAL, LITERAL_EXT)) { + if (resolvesToQualifiedNames(subscriptionExpr.getOperand(), context, LITERAL, LITERAL_EXT)) { return Optional .ofNullable(subscriptionExpr.getIndexExpression()) .map(index -> PyLiteralType.Companion.fromLiteralParameter(index, context.getTypeContext())) @@ -1285,6 +1293,13 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext< 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 qName = referenceExpression.getQualifier().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); @@ -2307,10 +2322,14 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext< } 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); + customContext = new Context(context, useFqn); context.getProcessingContext().put(TYPE_HINT_EVAL_CONTEXT, customContext); } try { @@ -2328,9 +2347,17 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext< private final @NotNull TypeEvalContext myContext; private final @NotNull Stack myTypeAliasStack = new Stack<>(); 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(); } diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyTypeParser.java b/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyTypeParser.java index 60b915dc4c20..6607eea69917 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyTypeParser.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyTypeParser.java @@ -124,6 +124,21 @@ public final class PyTypeParser { return parse(anchor, type, context).getType(); } + /** + * @param anchor should never be null or null will be returned + * @param context type evaluation context + * @param fqnOnly if true, resolves names using only fully-qualified lookup (ignores local scope/imported aliases) + * @return null either if there was an error during parsing or if extracted type is equivalent to Any or undefined + */ + @Contract("null, _, _, _ -> null") + public static @Nullable PyType getTypeByName(@Nullable PsiElement anchor, + @NotNull String type, + @NotNull TypeEvalContext context, + boolean fqnOnly) { + if (anchor == null) return EMPTY_RESULT.getType(); + return parse(anchor, type, context, fqnOnly).getType(); + } + /** * @param anchor should never be null or {@link PyTypeParser#EMPTY_RESULT} will be returned * @param type representation of the type to parse @@ -138,13 +153,23 @@ public final class PyTypeParser { * @param context type evaluation context */ public static @NotNull ParseResult parse(@NotNull PsiElement anchor, @NotNull String type, @NotNull TypeEvalContext context) { + return parse(anchor, type, context, false); + } + + /** + * @param anchor should never be null or {@link PyTypeParser#EMPTY_RESULT} will be returned + * @param type representation of the type to parse + * @param context type evaluation context + * @param fqnOnly if true, resolves names using only fully-qualified lookup (ignores local scope/imported aliases) + */ + public static @NotNull ParseResult parse(@NotNull PsiElement anchor, @NotNull String type, @NotNull TypeEvalContext context, boolean fqnOnly) { PyPsiUtils.assertValid(anchor); final ForwardDeclaration typeExpr = ForwardDeclaration.create(); final FunctionalParser classType = token(IDENTIFIER).then(many(op(".").skipThen(token(IDENTIFIER)))) - .map(new MakeSimpleType(anchor, context)) + .map(new MakeSimpleType(anchor, context, fqnOnly)) .cached() .named("class-type"); @@ -315,10 +340,16 @@ public final class PyTypeParser { private static class MakeSimpleType implements Function, List>>, ParseResult> { private final @NotNull PsiElement myAnchor; private final @NotNull TypeEvalContext myContext; + private final boolean myFqnOnly; MakeSimpleType(@NotNull PsiElement anchor, @NotNull TypeEvalContext context) { + this(anchor, context, false); + } + + MakeSimpleType(@NotNull PsiElement anchor, @NotNull TypeEvalContext context, boolean fqnOnly) { myAnchor = anchor; myContext = context; + myFqnOnly = fqnOnly; } @Override @@ -396,7 +427,24 @@ public final class PyTypeParser { final Token firstToken = tokens.get(0); final String firstText = firstToken.getText().toString(); final TextRange firstRange = firstToken.getRange(); - final List resolveResults = file.multiResolveName(firstText); + + final List resolveResults; + if (myFqnOnly) { + // First, try to resolve from "typing" for unqualified names (e.g., Literal, Any) + final var qNameContext = PyResolveImportUtil.fromFoothold(myAnchor); + final PsiElement typingMember = PyResolveImportUtil.resolveTopLevelMember(QualifiedName.fromDottedString("typing." + firstText), qNameContext); + if (typingMember != null) { + resolveResults = Collections.singletonList(new RatedResolveResult(RatedResolveResult.RATE_NORMAL, typingMember)); + } + else { + // Fall back to fully qualified name search (handled below using getImplicitlyResolvedType) + resolveResults = Collections.emptyList(); + } + } + else { + resolveResults = file.multiResolveName(firstText); + } + if (resolveResults.isEmpty()) { return getImplicitlyResolvedType(tokens, context, types, fullRanges, firstRange); }