From 157e636beaf6c3301e08b8ec55dba63a0794002d Mon Sep 17 00:00:00 2001
From: "Andrey.Matveev" <andrey.matveev@jetbrains.com>
Date: Thu, 28 Mar 2024 11:28:48 +0200
Subject: [PATCH] PY-71562 Make python support for
 FileIndexableEntitiesProvider

GitOrigin-RevId: 62f8c2b18d0d6b0988a85e4adf320f4200212fd0
---
 .idea/modules.xml                             |  1 +
 intellij.idea.community.main.iml              |  1 +
 .../build/CommunityRepositoryModules.kt       |  1 +
 .../resources/META-INF/plugin.xml             |  1 +
 ...ij.searchEverywhereMl.semantics.python.iml | 16 +++++++++
 ...ij.searchEverywhereMl.semantics.python.xml | 12 +++++++
 .../PythonIndexableEntitiesProvider.kt        | 34 +++++++++++++++++++
 7 files changed, 66 insertions(+)
 create mode 100644 plugins/search-everywhere-ml/semantics/python/intellij.searchEverywhereMl.semantics.python.iml
 create mode 100644 plugins/search-everywhere-ml/semantics/python/resources/intellij.searchEverywhereMl.semantics.python.xml
 create mode 100644 plugins/search-everywhere-ml/semantics/python/src/com.intellij.searchEverywhereMl.semantics.python/PythonIndexableEntitiesProvider.kt

diff --git a/.idea/modules.xml b/.idea/modules.xml
index dd1df8cd8363..4389cc3983cb 100644
--- a/.idea/modules.xml
+++ b/.idea/modules.xml
@@ -1183,6 +1183,7 @@
       <module fileurl="file://$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/intellij.searchEverywhereMl.semantics.iml" filepath="$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/intellij.searchEverywhereMl.semantics.iml" />
       <module fileurl="file://$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/java/intellij.searchEverywhereMl.semantics.java.iml" filepath="$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/java/intellij.searchEverywhereMl.semantics.java.iml" />
       <module fileurl="file://$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/kotlin/intellij.searchEverywhereMl.semantics.kotlin.iml" filepath="$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/kotlin/intellij.searchEverywhereMl.semantics.kotlin.iml" />
+      <module fileurl="file://$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/python/intellij.searchEverywhereMl.semantics.python.iml" filepath="$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/python/intellij.searchEverywhereMl.semantics.python.iml" />
       <module fileurl="file://$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/test-commands/intellij.searchEverywhereMl.semantics.testCommands.iml" filepath="$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/test-commands/intellij.searchEverywhereMl.semantics.testCommands.iml" />
       <module fileurl="file://$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/tests/intellij.searchEverywhereMl.semantics.tests.iml" filepath="$PROJECT_DIR$/plugins/search-everywhere-ml/semantics/tests/intellij.searchEverywhereMl.semantics.tests.iml" />
       <module fileurl="file://$PROJECT_DIR$/plugins/search-everywhere-ml/typos/intellij.searchEverywhereMl.typos.iml" filepath="$PROJECT_DIR$/plugins/search-everywhere-ml/typos/intellij.searchEverywhereMl.typos.iml" />
diff --git a/intellij.idea.community.main.iml b/intellij.idea.community.main.iml
index d7bbc16608c6..25270a5366e7 100644
--- a/intellij.idea.community.main.iml
+++ b/intellij.idea.community.main.iml
@@ -181,6 +181,7 @@
     <orderEntry type="module" module-name="intellij.searchEverywhereMl.typos" scope="TEST" />
     <orderEntry type="module" module-name="intellij.searchEverywhereMl.semantics.java" scope="RUNTIME" />
     <orderEntry type="module" module-name="intellij.searchEverywhereMl.semantics.kotlin" scope="RUNTIME" />
+    <orderEntry type="module" module-name="intellij.searchEverywhereMl.semantics.python" scope="RUNTIME" />
     <orderEntry type="module" module-name="intellij.toml" scope="RUNTIME" />
     <orderEntry type="module" module-name="intellij.platform.tracing.ide" scope="RUNTIME" />
     <orderEntry type="module" module-name="intellij.notebooks.visualization" scope="RUNTIME" />
diff --git a/platform/build-scripts/src/org/jetbrains/intellij/build/CommunityRepositoryModules.kt b/platform/build-scripts/src/org/jetbrains/intellij/build/CommunityRepositoryModules.kt
index 6256a30944e2..2df1028cc90e 100644
--- a/platform/build-scripts/src/org/jetbrains/intellij/build/CommunityRepositoryModules.kt
+++ b/platform/build-scripts/src/org/jetbrains/intellij/build/CommunityRepositoryModules.kt
@@ -292,6 +292,7 @@ object CommunityRepositoryModules {
     )) { spec ->
       spec.withModule("intellij.searchEverywhereMl.semantics.java")
       spec.withModule("intellij.searchEverywhereMl.semantics.kotlin")
+      spec.withModule("intellij.searchEverywhereMl.semantics.python")
       spec.withModule("intellij.searchEverywhereMl.semantics.testCommands")
     },
     plugin("intellij.platform.testFramework.ui") { spec ->
diff --git a/plugins/search-everywhere-ml/resources/META-INF/plugin.xml b/plugins/search-everywhere-ml/resources/META-INF/plugin.xml
index b33dc6952251..3b3df21f9322 100644
--- a/plugins/search-everywhere-ml/resources/META-INF/plugin.xml
+++ b/plugins/search-everywhere-ml/resources/META-INF/plugin.xml
@@ -22,6 +22,7 @@
     <module name="intellij.searchEverywhereMl.semantics"/>
     <module name="intellij.searchEverywhereMl.semantics.java"/>
     <module name="intellij.searchEverywhereMl.semantics.kotlin"/>
+    <module name="intellij.searchEverywhereMl.semantics.python"/>
     <module name="intellij.searchEverywhereMl.semantics.testCommands"/>
   </content>
 
diff --git a/plugins/search-everywhere-ml/semantics/python/intellij.searchEverywhereMl.semantics.python.iml b/plugins/search-everywhere-ml/semantics/python/intellij.searchEverywhereMl.semantics.python.iml
new file mode 100644
index 000000000000..84713a6ae98e
--- /dev/null
+++ b/plugins/search-everywhere-ml/semantics/python/intellij.searchEverywhereMl.semantics.python.iml
@@ -0,0 +1,16 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="JAVA_MODULE" version="4">
+  <component name="NewModuleRootManager" inherit-compiler-output="true">
+    <exclude-output />
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/resources" type="java-resource" />
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+    <orderEntry type="library" name="kotlin-stdlib" level="project" />
+    <orderEntry type="module" module-name="intellij.platform.ml.embeddings" />
+    <orderEntry type="module" module-name="intellij.python.psi" />
+    <orderEntry type="module" module-name="intellij.platform.core" />
+  </component>
+</module>
\ No newline at end of file
diff --git a/plugins/search-everywhere-ml/semantics/python/resources/intellij.searchEverywhereMl.semantics.python.xml b/plugins/search-everywhere-ml/semantics/python/resources/intellij.searchEverywhereMl.semantics.python.xml
new file mode 100644
index 000000000000..7cf1286ef5b5
--- /dev/null
+++ b/plugins/search-everywhere-ml/semantics/python/resources/intellij.searchEverywhereMl.semantics.python.xml
@@ -0,0 +1,12 @@
+<idea-plugin package="com.intellij.searchEverywhereMl.semantics.python">
+  <dependencies>
+    <plugin id="com.intellij.modules.python"/>
+    <module name="intellij.python.community.impl"/>
+    <module name="intellij.searchEverywhereMl.semantics"/>
+  </dependencies>
+
+  <extensions defaultExtensionNs="com.intellij.searcheverywhere.ml">
+    <fileIndexableEntitiesProvider
+      implementation="com.intellij.searchEverywhereMl.semantics.python.PythonIndexableEntitiesProvider"/>
+  </extensions>
+</idea-plugin>
\ No newline at end of file
diff --git a/plugins/search-everywhere-ml/semantics/python/src/com.intellij.searchEverywhereMl.semantics.python/PythonIndexableEntitiesProvider.kt b/plugins/search-everywhere-ml/semantics/python/src/com.intellij.searchEverywhereMl.semantics.python/PythonIndexableEntitiesProvider.kt
new file mode 100644
index 000000000000..3e06e5c4a947
--- /dev/null
+++ b/plugins/search-everywhere-ml/semantics/python/src/com.intellij.searchEverywhereMl.semantics.python/PythonIndexableEntitiesProvider.kt
@@ -0,0 +1,34 @@
+package com.intellij.searchEverywhereMl.semantics.python
+
+import com.intellij.platform.ml.embeddings.search.indices.FileIndexableEntitiesProvider
+import com.intellij.platform.ml.embeddings.search.services.IndexableClass
+import com.intellij.platform.ml.embeddings.search.services.IndexableSymbol
+import com.intellij.psi.PsiFile
+import com.intellij.psi.util.PsiTreeUtil
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyFile
+import com.jetbrains.python.psi.PyFunction
+
+class PythonIndexableEntitiesProvider : FileIndexableEntitiesProvider {
+  override fun extractIndexableSymbols(file: PsiFile): List<IndexableSymbol> {
+    return when (file) {
+      is PyFile -> PsiTreeUtil.findChildrenOfAnyType(file, false, PyFunction::class.java)
+        .filter { it.name != ANONYMOUS_ID }
+        .map { IndexableSymbol(it.name?.intern() ?: "") }
+      else -> emptyList()
+    }
+  }
+
+  override fun extractIndexableClasses(file: PsiFile): List<IndexableClass> {
+    return when (file) {
+      is PyFile -> PsiTreeUtil.getStubChildrenOfTypeAsList(file, PyClass::class.java)
+        .filter { it.name != ANONYMOUS_ID }
+        .map { IndexableClass(it.name?.intern() ?: "") }
+      else -> emptyList()
+    }
+  }
+
+  companion object {
+    private const val ANONYMOUS_ID = "<anonymous>"
+  }
+}
\ No newline at end of file