Skip to content

Commit 7a3f724

Browse files
Will-ShaoHuaWill-ShaoHua
authored
[CodeWhisperer] Add cwspr support for JSX file (#3296)
* Add cwspr jsx support * Improve CodeWhispererLanguageManager
1 parent b429407 commit 7a3f724

File tree

9 files changed

+324
-42
lines changed

9 files changed

+324
-42
lines changed

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorListener.kt‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,26 +9,27 @@ import com.intellij.openapi.editor.event.EditorFactoryEvent
99
import com.intellij.openapi.editor.event.EditorFactoryListener
1010
import com.intellij.openapi.editor.impl.EditorImpl
1111
import com.intellij.psi.PsiDocumentManager
12-
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.codeWhispererLanguage
12+
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.programmingLanguage
1313
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
1414
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererLanguageManager
15+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.toCodeWhispererLanguage
1516
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus
1617
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererCodeCoverageTracker
1718

1819
class CodeWhispererEditorListener : EditorFactoryListener {
1920
override fun editorCreated(event: EditorFactoryEvent) {
2021
val editor = (event.editor as? EditorImpl) ?: return
2122
editor.project?.let { project ->
22-
PsiDocumentManager.getInstance(project).getPsiFile(editor.document)?.codeWhispererLanguage ?. let { language ->
23+
PsiDocumentManager.getInstance(project).getPsiFile(editor.document)?.programmingLanguage ?. let { language ->
2324
// If language is not supported by CodeWhisperer, no action needed
24-
if (!CodeWhispererLanguageManager.getInstance().isLanguageSupported(language.toString())) return
25+
if (!CodeWhispererLanguageManager.getInstance().isLanguageSupported(language)) return
2526
// If language is supported, install document listener for CodeWhisperer service
2627
editor.document.addDocumentListener(
2728
object : DocumentListener {
2829
override fun documentChanged(event: DocumentEvent) {
2930
if (!CodeWhispererExplorerActionManager.getInstance().hasAcceptedTermsOfService()) return
3031
CodeWhispererInvocationStatus.getInstance().documentChanged()
31-
CodeWhispererCodeCoverageTracker.getInstance(language).apply {
32+
CodeWhispererCodeCoverageTracker.getInstance(language.toCodeWhispererLanguage()).apply {
3233
activateTrackerIfNotActive()
3334
documentChanged(event)
3435
}

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorUtil.kt‎

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import com.intellij.openapi.vfs.VfsUtilCore
1313
import com.intellij.openapi.vfs.VirtualFile
1414
import com.intellij.psi.PsiFile
1515
import com.intellij.ui.popup.AbstractPopup
16-
import software.amazon.awssdk.utils.StringUtils.lowerCase
16+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.toCodeWhispererLanguage
1717
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretContext
1818
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition
1919
import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
@@ -36,24 +36,19 @@ object CodeWhispererEditorUtil {
3636
return FileContextInfo(caretContext, fileName, programmingLanguage)
3737
}
3838

39-
fun ProgrammingLanguage.toCodeWhispererLanguage() = when (languageName) {
40-
CodewhispererLanguage.Python.toString() -> CodewhispererLanguage.Python
41-
CodewhispererLanguage.Java.toString() -> CodewhispererLanguage.Java
42-
CodewhispererLanguage.Javascript.toString() -> CodewhispererLanguage.Javascript
43-
"plain_text" -> CodewhispererLanguage.Plaintext
44-
else -> CodewhispererLanguage.Unknown
45-
}
46-
39+
/**
40+
* return IDE identifiable ProgrammingLanguage e.g. ProgrammingLanguage("java"), ProgrammingLanguage("jsx harmony") etc.
41+
*/
4742
val PsiFile.programmingLanguage: ProgrammingLanguage
48-
get() = ProgrammingLanguage(lowerCase(this.fileType.name))
43+
get() = ProgrammingLanguage(this.fileType.name)
4944

5045
val PsiFile.codeWhispererLanguage: CodewhispererLanguage
5146
get() = this.programmingLanguage.toCodeWhispererLanguage()
5247

5348
val VirtualFile.codeWhispererLanguage: CodewhispererLanguage
54-
get() = ProgrammingLanguage(fileType.name.lowercase()).toCodeWhispererLanguage()
49+
get() = ProgrammingLanguage(fileType.name).toCodeWhispererLanguage()
5550

56-
private fun extractCaretContext(editor: Editor): CaretContext {
51+
fun extractCaretContext(editor: Editor): CaretContext {
5752
val document = editor.document
5853
val caretOffset = editor.caretModel.primaryCaret.offset
5954
val totalCharLength = editor.document.textLength

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/CodeWhispererLanguageManager.kt‎

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,43 @@
44
package software.aws.toolkits.jetbrains.services.codewhisperer.language
55

66
import com.intellij.openapi.components.service
7+
import software.aws.toolkits.jetbrains.services.codewhisperer.model.ProgrammingLanguage
78
import software.aws.toolkits.telemetry.CodewhispererLanguage
89

910
class CodeWhispererLanguageManager {
10-
fun isLanguageSupported(language: String): Boolean =
11-
language == CodewhispererLanguage.Java.toString() ||
12-
language == CodewhispererLanguage.Python.toString() ||
13-
language == CodewhispererLanguage.Javascript.toString()
11+
private val supportedLanguage = setOf(
12+
CodewhispererLanguage.Java.toString(),
13+
CodewhispererLanguage.Python.toString(),
14+
CodewhispererLanguage.Javascript.toString()
15+
)
16+
17+
fun isLanguageSupported(language: ProgrammingLanguage): Boolean {
18+
val mappedLanguage = getParentLanguage(language)
19+
return supportedLanguage.contains(mappedLanguage.languageName)
20+
}
21+
22+
/**
23+
* This should be called to map some language dialect to their mother language
24+
* e.g. JSX -> JavaScript, TypeScript -> JavaScript etc.
25+
*/
26+
internal fun getParentLanguage(language: ProgrammingLanguage): ProgrammingLanguage =
27+
when {
28+
language.languageName.contains("jsx") -> ProgrammingLanguage(CodewhispererLanguage.Javascript)
29+
else -> language
30+
}
1431

1532
companion object {
1633
fun getInstance(): CodeWhispererLanguageManager = service()
1734
}
1835
}
36+
37+
fun ProgrammingLanguage.toCodeWhispererLanguage() = when {
38+
languageName.contains("python") -> CodewhispererLanguage.Python
39+
languageName.contains("javascript") -> CodewhispererLanguage.Javascript
40+
languageName.contains("java") -> CodewhispererLanguage.Java
41+
languageName.contains("jsx") -> CodewhispererLanguage.Jsx
42+
languageName.contains("plain_text") -> CodewhispererLanguage.Plaintext
43+
else -> CodewhispererLanguage.Unknown
44+
}
45+
46+
fun CodewhispererLanguage.toProgrammingLanguage() = ProgrammingLanguage(this.toString())

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionco
1212
import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContext
1313
import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext
1414
import software.aws.toolkits.telemetry.CodewhispererAutomatedTriggerType
15+
import software.aws.toolkits.telemetry.CodewhispererLanguage
1516
import software.aws.toolkits.telemetry.CodewhispererTriggerType
1617
import software.aws.toolkits.telemetry.Result
1718

@@ -23,7 +24,18 @@ data class FileContextInfo(
2324
val programmingLanguage: ProgrammingLanguage
2425
)
2526

26-
data class ProgrammingLanguage(val languageName: String)
27+
data class ProgrammingLanguage(private val myLanguageName: String) {
28+
// case-insensitive
29+
val languageName get() = myLanguageName.lowercase()
30+
constructor(codewhispererLanguage: CodewhispererLanguage) : this(codewhispererLanguage.toString())
31+
32+
override fun equals(other: Any?): Boolean {
33+
if (other !is ProgrammingLanguage) return false
34+
return other.languageName == this.languageName
35+
}
36+
37+
override fun toString(): String = languageName
38+
}
2739

2840
data class RecommendationContext(
2941
val details: List<DetailContext>,

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt‎

Lines changed: 19 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ class CodeWhispererService {
8484
return
8585
}
8686

87-
val language = requestContext.fileContextInfo.programmingLanguage.languageName
87+
val language = requestContext.fileContextInfo.programmingLanguage
8888
if (!CodeWhispererLanguageManager.getInstance().isLanguageSupported(language)) {
8989
LOG.debug { "Programming language $language is not supported by CodeWhisperer" }
9090
if (triggerTypeInfo.triggerType == CodewhispererTriggerType.OnDemand) {
@@ -474,24 +474,6 @@ class CodeWhispererService {
474474
return response.toBuilder().recommendations(validatedRecommendations).build()
475475
}
476476

477-
private fun buildCodeWhispererRequest(
478-
fileContextInfo: FileContextInfo
479-
): ListRecommendationsRequest {
480-
val programmingLanguage = ProgrammingLanguage.builder()
481-
.languageName(fileContextInfo.programmingLanguage.languageName)
482-
.build()
483-
val fileContext = FileContext.builder()
484-
.leftFileContent(fileContextInfo.caretContext.leftFileContext)
485-
.rightFileContent(fileContextInfo.caretContext.rightFileContext)
486-
.filename(fileContextInfo.filename)
487-
.programmingLanguage(programmingLanguage)
488-
.build()
489-
490-
return ListRecommendationsRequest.builder()
491-
.fileContext(fileContext)
492-
.build()
493-
}
494-
495477
private fun buildInvocationContext(
496478
requestContext: RequestContext,
497479
responseContext: ResponseContext,
@@ -585,6 +567,24 @@ class CodeWhispererService {
585567
val KEY_CODEWHISPERER_METADATA: Key<CodeWhispererMetadata> = Key.create("codewhisperer.metadata")
586568
fun getInstance(): CodeWhispererService = service()
587569
const val KET_SESSION_ID = "x-amzn-SessionId"
570+
571+
internal fun buildCodeWhispererRequest(
572+
fileContextInfo: FileContextInfo
573+
): ListRecommendationsRequest {
574+
val programmingLanguage = ProgrammingLanguage.builder()
575+
.languageName(CodeWhispererLanguageManager.getInstance().getParentLanguage(fileContextInfo.programmingLanguage).languageName)
576+
.build()
577+
val fileContext = FileContext.builder()
578+
.leftFileContent(fileContextInfo.caretContext.leftFileContext)
579+
.rightFileContent(fileContextInfo.caretContext.rightFileContext)
580+
.filename(fileContextInfo.filename)
581+
.programmingLanguage(programmingLanguage)
582+
.build()
583+
584+
return ListRecommendationsRequest.builder()
585+
.fileContext(fileContext)
586+
.build()
587+
}
588588
}
589589
}
590590

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererCodeCoverageTracker.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import info.debatty.java.stringsimilarity.Levenshtein
1717
import org.jetbrains.annotations.TestOnly
1818
import software.aws.toolkits.core.utils.debug
1919
import software.aws.toolkits.core.utils.getLogger
20-
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.toCodeWhispererLanguage
20+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.toCodeWhispererLanguage
2121
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
2222
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
2323
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import com.intellij.openapi.vfs.VirtualFile
1111
import software.amazon.awssdk.services.codewhisperer.model.Recommendation
1212
import software.aws.toolkits.core.utils.debug
1313
import software.aws.toolkits.core.utils.getLogger
14-
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.toCodeWhispererLanguage
14+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.toCodeWhispererLanguage
1515
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CodeScanTelemetryEvent
1616
import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationContext
1717
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext

‎jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererEditorUtilTest.kt‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.codewhisperer
5+
6+
import com.intellij.openapi.command.WriteCommandAction
7+
import com.intellij.openapi.fileTypes.FileType
8+
import com.intellij.openapi.project.Project
9+
import com.intellij.psi.PsiFile
10+
import com.intellij.testFramework.fixtures.CodeInsightTestFixture
11+
import com.intellij.testFramework.runInEdtAndGet
12+
import org.assertj.core.api.Assertions.assertThat
13+
import org.junit.Before
14+
import org.junit.Rule
15+
import org.junit.Test
16+
import org.mockito.kotlin.doReturn
17+
import org.mockito.kotlin.mock
18+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName
19+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
20+
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil
21+
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.programmingLanguage
22+
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
23+
24+
class CodeWhispererEditorUtilTest {
25+
@Rule
26+
@JvmField
27+
var projectRule = PythonCodeInsightTestFixtureRule()
28+
29+
private lateinit var project: Project
30+
private lateinit var fixture: CodeInsightTestFixture
31+
32+
@Before
33+
fun setup() {
34+
project = projectRule.project
35+
fixture = projectRule.fixture
36+
}
37+
38+
@Test
39+
fun `test psiFile_programmingLanguage`() {
40+
val fileTypeMock = mock<FileType> {
41+
on { name } doReturn "Java"
42+
}
43+
val psiFile = mock<PsiFile> {
44+
on { fileType } doReturn fileTypeMock
45+
}
46+
val programmingLanguage = psiFile.programmingLanguage
47+
assertThat(programmingLanguage.languageName).isEqualTo("java")
48+
}
49+
50+
@Test
51+
fun `test getFileContextInfo`() {
52+
val psiFile = fixture.configureByText(pythonFileName, pythonTestLeftContext)
53+
val fileContext = runInEdtAndGet {
54+
fixture.editor.caretModel.moveToOffset(fixture.editor.document.textLength)
55+
CodeWhispererEditorUtil.getFileContextInfo(fixture.editor, psiFile)
56+
}
57+
58+
assertThat(fileContext.filename).isEqualTo(pythonFileName)
59+
assertThat(fileContext.programmingLanguage.languageName).isEqualTo("python")
60+
assertThat(fileContext.caretContext.leftFileContext).isEqualTo(pythonTestLeftContext)
61+
assertThat(fileContext.caretContext.rightFileContext).isEqualTo("")
62+
}
63+
64+
/**
65+
* # function to add 2 numbers
66+
* def addTwoNumbers(a, b)
67+
* ^
68+
* cursor
69+
*/
70+
@Test
71+
fun `test extractCaretContext`() {
72+
val pythonComment = "# function to add 2 numbers\n"
73+
val pythonTestRightContext = "(a, b)"
74+
fixture.configureByText(pythonFileName, pythonComment)
75+
val caretContext = runInEdtAndGet {
76+
WriteCommandAction.runWriteCommandAction(project) {
77+
fixture.editor.caretModel.moveToOffset(fixture.editor.document.textLength)
78+
fixture.editor.document.insertString(fixture.editor.caretModel.offset, pythonTestLeftContext)
79+
80+
fixture.editor.caretModel.moveToOffset(fixture.editor.document.textLength)
81+
fixture.editor.document.insertString(fixture.editor.caretModel.offset, pythonTestRightContext)
82+
}
83+
84+
CodeWhispererEditorUtil.extractCaretContext(fixture.editor)
85+
}
86+
87+
assertThat(caretContext.leftFileContext).isEqualTo("$pythonComment$pythonTestLeftContext")
88+
assertThat(caretContext.rightFileContext).isEqualTo(pythonTestRightContext)
89+
assertThat(caretContext.leftContextOnCurrentLine).isEqualTo(pythonTestLeftContext)
90+
}
91+
}

0 commit comments

Comments
 (0)