crossfile/utg patch (#3768)

Will-ShaoHua · web-flow · commit 19742b80c6c8 · 2023-07-19T02:36:38.000-07:00
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererJsx.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererJsx.kt
@@ -19,6 +19,8 @@ class CodeWhispererJsx private constructor() : CodeWhispererProgrammingLanguage(
 
     override fun isAllClassifier(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "jsx"
 
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererTsx.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererTsx.kt
@@ -19,6 +19,8 @@ class CodeWhispererTsx private constructor() : CodeWhispererProgrammingLanguage(
 
     override fun isAllClassifier(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "tsx"
 
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt
@@ -25,7 +25,9 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhisper
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJavaScript
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJsx
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTsx
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTypeScript
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.Chunk
@@ -51,8 +53,8 @@ private val codewhispererCodeChunksIndex = GistManager.getInstance()
 private fun getFileCrawlerForLanguage(programmingLanguage: CodeWhispererProgrammingLanguage) = when (programmingLanguage) {
     is CodeWhispererJava -> JavaCodeWhispererFileCrawler
     is CodeWhispererPython -> PythonCodeWhispererFileCrawler
-    is CodeWhispererJavaScript -> JavascriptCodeWhispererFileCrawler
-    is CodeWhispererTypeScript -> TypescriptCodeWhispererFileCrawler
+    is CodeWhispererJavaScript, is CodeWhispererJsx -> JavascriptCodeWhispererFileCrawler
+    is CodeWhispererTypeScript, is CodeWhispererTsx -> TypescriptCodeWhispererFileCrawler
     else -> NoOpFileCrawler()
 }
 
@@ -127,11 +129,12 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
             when (language) {
                 is CodeWhispererJava -> extractSupplementalFileContextForSrc(psiFile, targetContext)
 
-                is CodeWhispererPython, is CodeWhispererJavaScript, is CodeWhispererTypeScript -> if (userGroup == CodeWhispererUserGroup.CrossFile) {
-                    extractSupplementalFileContextForSrc(psiFile, targetContext)
-                } else {
-                    emptyList()
-                }
+                is CodeWhispererPython, is CodeWhispererJavaScript, is CodeWhispererTypeScript, is CodeWhispererJsx, is CodeWhispererTsx ->
+                    if (userGroup == CodeWhispererUserGroup.CrossFile) {
+                        extractSupplementalFileContextForSrc(psiFile, targetContext)
+                    } else {
+                        emptyList()
+                    }
 
                 else -> emptyList()
             }
@@ -195,10 +198,12 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
         return chunks.take(CodeWhispererConstants.CrossFile.CHUNK_SIZE)
     }
 
-    override fun isTestFile(psiFile: PsiFile) = when (psiFile.programmingLanguage()) {
-        is CodeWhispererJava -> TestSourcesFilter.isTestSources(psiFile.virtualFile, project)
-        is CodeWhispererPython -> PythonCodeWhispererFileCrawler.testFilenamePattern.matches(psiFile.name)
-        else -> true
+    override fun isTestFile(psiFile: PsiFile): Boolean {
+        val path = runReadAction { contentRootPathProvider.getPathToElement(project, psiFile.virtualFile, null) ?: psiFile.virtualFile.path }
+        return TestSourcesFilter.isTestSources(psiFile.virtualFile, project) ||
+            path.contains("""test/""") ||
+            path.contains("""tst/""") ||
+            path.contains("""tests/""")
     }
 
     @VisibleForTesting
@@ -255,23 +260,25 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
         val focalFile = getFileCrawlerForLanguage(targetContext.programmingLanguage).findFocalFileForTest(psiFile)
 
         return focalFile?.let { file ->
-            val relativePath = contentRootPathProvider.getPathToElement(project, file, null) ?: file.path
-            val content = file.content()
+            runReadAction {
+                val relativePath = contentRootPathProvider.getPathToElement(project, file, null) ?: file.path
+                val content = file.content()
 
-            if (content.isBlank()) {
-                emptyList()
-            } else {
-                listOf(
-                    Chunk(
-                        content = CodeWhispererConstants.Utg.UTG_PREFIX + file.content().let {
-                            it.substring(
-                                0,
-                                minOf(it.length, CodeWhispererConstants.Utg.UTG_SEGMENT_SIZE)
-                            )
-                        },
-                        path = relativePath
+                if (content.isBlank()) {
+                    emptyList()
+                } else {
+                    listOf(
+                        Chunk(
+                            content = CodeWhispererConstants.Utg.UTG_PREFIX + file.content().let {
+                                it.substring(
+                                    0,
+                                    minOf(it.length, CodeWhispererConstants.Utg.UTG_SEGMENT_SIZE)
+                                )
+                            },
+                            path = relativePath
+                        )
                     )
-                )
+                }
             }
         }.orEmpty()
     }
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt
@@ -12,6 +12,8 @@ import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiManager
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
 
 /**
  * An interface define how do we parse and fetch files provided a psi file or project
@@ -70,12 +72,19 @@ abstract class CodeWhispererFileCrawler : FileCrawler {
 
     override fun listRelevantFilesInEditors(psiFile: PsiFile): List<VirtualFile> {
         val targetFile = psiFile.virtualFile
+        val language = psiFile.programmingLanguage()
+
+        val isTestFilePredicate: (file: VirtualFile, project: Project) -> Boolean = if (language is CodeWhispererJava) {
+            { file, project -> TestSourcesFilter.isTestSources(file, project) }
+        } else {
+            { file, _ -> testFilenamePattern.matches(file.name) }
+        }
 
         val openedFiles = runReadAction {
             FileEditorManager.getInstance(psiFile.project).openFiles.toList().filter {
                 it.name != psiFile.virtualFile.name &&
                     isSameDialect(psiFile.virtualFile.extension) &&
-                    !TestSourcesFilter.isTestSources(it, psiFile.project)
+                    !isTestFilePredicate(it, psiFile.project)
             }
         }
 
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavascriptCodeWhispererFileCrawler.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavascriptCodeWhispererFileCrawler.kt
@@ -10,8 +10,7 @@ object JavascriptCodeWhispererFileCrawler : CodeWhispererFileCrawler() {
     override val fileExtension: String = "js"
     override val dialects: Set<String> = setOf("js", "jsx")
 
-    // TODO: Add implementation when UTG is enabled
-    override val testFilenamePattern: Regex = "".toRegex()
+    override val testFilenamePattern: Regex = """^.*\.test\.(js|jsx)${'$'}""".toRegex()
 
     // TODO: Add implementation when UTG is enabled
     override fun guessSourceFileName(tstFileName: String): String = ""
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/PythonCodeWhispererFileCrawler.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/PythonCodeWhispererFileCrawler.kt
@@ -4,13 +4,9 @@
 package software.aws.toolkits.jetbrains.services.codewhisperer.util
 
 import com.intellij.openapi.application.runReadAction
-import com.intellij.openapi.module.ModuleUtilCore
-import com.intellij.openapi.project.rootManager
-import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiFile
 import com.jetbrains.python.psi.PyFile
-import org.jetbrains.jps.model.java.JavaModuleSourceRootTypes
 
 object PythonCodeWhispererFileCrawler : CodeWhispererFileCrawler() {
     override val fileExtension: String = "py"
@@ -36,15 +32,12 @@ object PythonCodeWhispererFileCrawler : CodeWhispererFileCrawler() {
 
     override fun findFocalFileForTest(psiFile: PsiFile): VirtualFile? = findSourceFileByName(psiFile) ?: findRelevantFileFromEditors(psiFile)
 
-    private fun findSourceFileByName(psiFile: PsiFile): VirtualFile? {
-        val module = ModuleUtilCore.findModuleForFile(psiFile)
-
-        return module?.rootManager?.getSourceRoots(JavaModuleSourceRootTypes.PRODUCTION)?.let { srcRoot ->
-            srcRoot
-                .map { root -> VfsUtil.collectChildrenRecursively(root) }
-                .flatten()
-                .find { !it.isDirectory && it.isWritable && it.name.contains(guessSourceFileName(psiFile.name)) }
-        }
+    private fun findSourceFileByName(psiFile: PsiFile): VirtualFile? = super.listFilesUnderProjectRoot(psiFile.project).find {
+        !it.isDirectory &&
+            it.isWritable &&
+            it.name != psiFile.virtualFile.name &&
+            // TODO: should we use strict equal instead?
+            it.name.contains(guessSourceFileName(psiFile.name))
     }
 
     /**
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/TypescriptCodeWhispererFileCrawler.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/TypescriptCodeWhispererFileCrawler.kt
@@ -10,8 +10,7 @@ object TypescriptCodeWhispererFileCrawler : CodeWhispererFileCrawler() {
     override val fileExtension: String = "ts"
     override val dialects: Set<String> = setOf("ts", "tsx")
 
-    // TODO: Add implementation when UTG is enabled
-    override val testFilenamePattern: Regex = "".toRegex()
+    override val testFilenamePattern: Regex = """^.*\.test\.(ts|tsx)${'$'}""".toRegex()
 
     // TODO: Add implementation when UTG is enabled
     override fun guessSourceFileName(tstFileName: String): String = ""
