[python] PY-83505 ExternalTypeProvider extension point

GitOrigin-RevId: 4c576007a0f84542f5b91f5e502d0e280c3d5b7c

python/python-psi-api/resources/intellij.python.psi.xml

@@ -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>

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<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));
   }
@@ -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);

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<PyType?>?
+
+  companion object {
+    @JvmField
+    val EP_NAME: ExtensionPointName<TypeEvalExternalTypeProvider> =
+      ExtensionPointName.create("Pythonid.typeEvalExternalTypeProvider")
+  }
+}