From 8b387672f296b8dc83fb45b9285601963f0652cb Mon Sep 17 00:00:00 2001
From: Andrey Vokin <andrey.vokin@jetbrains.com>
Date: Mon, 23 Jun 2025 10:08:25 +0200
Subject: [PATCH] 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
---
 .../resources/intellij.python.psi.impl.xml    |  1 +
 .../python/psi/PyFileElementType.java         |  2 +-
 .../stubs/PyImportElementElementType.java     | 14 +++++++---
 .../python/psi/stubs/PyImportNameIndex.kt     | 27 +++++++++++++++++++
 4 files changed, 39 insertions(+), 5 deletions(-)
 create mode 100644 python/python-psi-impl/src/com/jetbrains/python/psi/stubs/PyImportNameIndex.kt

diff --git a/python/python-psi-impl/resources/intellij.python.psi.impl.xml b/python/python-psi-impl/resources/intellij.python.psi.impl.xml
index 2d48452f760c..661b25ab35e5 100644
--- a/python/python-psi-impl/resources/intellij.python.psi.impl.xml
+++ b/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"/>
 
diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java b/python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java
index cc71b00ef9f5..a7d4e4ee7aeb 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java
+++ b/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
diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyImportElementElementType.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyImportElementElementType.java
index 92108ba841b3..9581b15c73b4 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyImportElementElementType.java
+++ b/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);
diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/stubs/PyImportNameIndex.kt b/python/python-psi-impl/src/com/jetbrains/python/psi/stubs/PyImportNameIndex.kt
new file mode 100644
index 000000000000..320fb6feec1a
--- /dev/null
+++ b/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
+}
\ No newline at end of file