[python] PY-83505 ExternalTypeProvider extension point

(cherry picked from commit 4c576007a0f84542f5b91f5e502d0e280c3d5b7c)

GitOrigin-RevId: 57ce4325e4948c00c38a8d362c97b89de3f8b854
This commit is contained in:
Morgan Bartholomew
2025-11-05 02:02:56 +00:00
committed by intellij-monorepo-bot
parent 304d13f23e
commit 2f20b95624
5 changed files with 141 additions and 4 deletions
@@ -41,5 +41,8 @@
<extensionPoint qualifiedName="Pythonid.pythonDocumentationQuickInfoProvider"
interface="com.jetbrains.python.documentation.PythonDocumentationQuickInfoProvider"
dynamic="true"/>
<extensionPoint qualifiedName="Pythonid.typeEvalExternalTypeProvider"
interface="com.jetbrains.python.psi.types.TypeEvalExternalTypeProvider"
dynamic="true"/>
</extensionPoints>
</idea-plugin>
@@ -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<ProcessingContext> myProcessingContext = ThreadLocal.withInitial(ProcessingContext::new);
protected final Map<PyTypedElement, PyType> myEvaluated = CollectionFactory.createConcurrentSoftValueMap();
public final Map<PyTypedElement, PyType> myExternalEvaluated = CollectionFactory.createConcurrentSoftValueMap();
protected final Map<PyCallable, PyType> myEvaluatedReturn = CollectionFactory.createConcurrentSoftValueMap();
protected final Map<Pair<PyExpression, Object>, 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));
}
@@ -213,6 +218,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);
@@ -222,6 +232,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;
}
@@ -234,6 +249,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();
@@ -270,6 +290,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);
@@ -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<PyType?>?
companion object {
@JvmField
val EP_NAME: ExtensionPointName<TypeEvalExternalTypeProvider> =
ExtensionPointName.create("Pythonid.typeEvalExternalTypeProvider")
}
}
@@ -783,6 +783,11 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext<
return false;
}
@ApiStatus.Internal
public static @Nullable Ref<PyType> getType(@NotNull PyExpression expression, @NotNull TypeEvalContext context, boolean useFqn) {
return staticWithCustomContext(context, useFqn, customContext -> getType(expression, customContext));
}
public static @Nullable Ref<PyType> 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<PyType> 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> T staticWithCustomContext(@NotNull TypeEvalContext context, @NotNull Function<@NotNull Context, T> delegate) {
return staticWithCustomContext(context, false, delegate);
}
private static <T> 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<PyQualifiedNameOwner> 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();
}
@@ -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 <tt>Any</tt> or <tt>undefined</tt>
*/
@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<ParseResult, PyElementType> typeExpr = ForwardDeclaration.create();
final FunctionalParser<ParseResult, PyElementType> 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<Pair<Token<PyElementType>, List<Token<PyElementType>>>, 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<PyElementType> firstToken = tokens.get(0);
final String firstText = firstToken.getText().toString();
final TextRange firstRange = firstToken.getRange();
final List<RatedResolveResult> resolveResults = file.multiResolveName(firstText);
final List<RatedResolveResult> 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);
}