PY-79510 PyCharm: Code entry points for AI playground

Add an index to the names imported from Python. This will allow you to check if the project has AI libraries in its dependencies.

GitOrigin-RevId: d16f3163a602194ef0f3254402dd4061ff8d2f7f

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

@@ -44,6 +44,7 @@
     <stubIndex implementation="com.jetbrains.python.psi.stubs.PyClassAttributesIndex"/>
     <stubIndex implementation="com.jetbrains.python.psi.stubs.PyExportedModuleAttributeIndex"/>
     <stubIndex implementation="com.jetbrains.python.psi.stubs.PyTypeAliasNameIndex"/>
+    <stubIndex implementation="com.jetbrains.python.psi.stubs.PyImportNameIndex"/>
     <fileBasedIndex implementation="com.jetbrains.python.psi.stubs.PyModuleNameIndex"/>
     <fileBasedIndex implementation="com.jetbrains.python.psi.stubs.PySetuptoolsNamespaceIndex"/>
 

python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java

@@ -60,7 +60,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
   @Override
   public int getStubVersion() {
     // Don't forget to update versions of indexes that use the updated stub-based elements
-    return 101;
+    return 102;
   }
 
   @Override

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyImportElementElementType.java

@@ -3,10 +3,7 @@ package com.jetbrains.python.psi.impl.stubs;
 
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.stubs.IStubElementType;
-import com.intellij.psi.stubs.StubElement;
-import com.intellij.psi.stubs.StubInputStream;
-import com.intellij.psi.stubs.StubOutputStream;
+import com.intellij.psi.stubs.*;
 import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.psi.PyImportElement;
@@ -14,6 +11,7 @@ import com.jetbrains.python.psi.PyStubElementType;
 import com.jetbrains.python.psi.PyTargetExpression;
 import com.jetbrains.python.psi.impl.PyImportElementImpl;
 import com.jetbrains.python.psi.stubs.PyImportElementStub;
+import com.jetbrains.python.psi.stubs.PyImportNameIndex;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 
@@ -45,6 +43,14 @@ public class PyImportElementElementType extends PyStubElementType<PyImportElemen
     return new PyImportElementStubImpl(psi.getImportedQName(), asName != null ? asName.getName() : "", parentStub, getStubElementType());
   }
 
+  @Override
+  public void indexStub(@NotNull PyImportElementStub stub, @NotNull IndexSink sink) {
+    QualifiedName name = stub.getImportedQName();
+    if (name != null) {
+      sink.occurrence(PyImportNameIndex.KEY, name.toString());
+    }
+  }
+
   @Override
   public void serialize(@NotNull PyImportElementStub stub, @NotNull StubOutputStream dataStream) throws IOException {
     QualifiedName.serialize(stub.getImportedQName(), dataStream);

python/python-psi-impl/src/com/jetbrains/python/psi/stubs/PyImportNameIndex.kt

@@ -0,0 +1,27 @@
+package com.jetbrains.python.psi.stubs
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.stubs.StringStubIndexExtension
+import com.intellij.psi.stubs.StubIndex
+import com.intellij.psi.stubs.StubIndexKey
+import com.jetbrains.python.psi.PyImportElement
+
+class PyImportNameIndex : StringStubIndexExtension<PyImportElement>() {
+  companion object {
+    @JvmField
+    val  KEY: StubIndexKey<String, PyImportElement> = StubIndexKey.createIndexKey("python.import.name")
+
+    fun containsImport(name: String, project: Project): Boolean {
+      val projectScope = GlobalSearchScope.projectScope(project)
+      var found = false
+      StubIndex.getInstance().processElements(KEY, name, project, projectScope, PyImportElement::class.java) {
+        found = true
+        return@processElements false
+      }
+      return found
+    }
+  }
+
+  override fun getKey(): StubIndexKey<String, PyImportElement> = KEY
+}