From 6786175c34ea9a50d8f69935c099c880d359355e Mon Sep 17 00:00:00 2001 From: Phodal Huang Date: Tue, 26 Dec 2023 22:10:00 +0800 Subject: [PATCH] feat(python): fix insert issues --- gradle-222.properties | 2 +- .../provider/PythonLivingDocumentation.kt | 71 ++++++++++--------- .../custom/document/CustomLivingDocTask.kt | 4 +- .../action/task/LivingDocPromptBuilder.kt | 4 +- 4 files changed, 42 insertions(+), 39 deletions(-) diff --git a/gradle-222.properties b/gradle-222.properties index 6dcbf936fb..c0e69c7646 100644 --- a/gradle-222.properties +++ b/gradle-222.properties @@ -4,5 +4,5 @@ ideaVersion=IU-2022.2.4 # please see https://www.jetbrains.org/intellij/sdk/docs/basics/getting_started/build_number_ranges.html for description -pluginSinceBuild = 222.* +pluginSinceBuild = 221.* pluginUntilBuild = 232.* diff --git a/pycharm/src/main/kotlin/cc/unitmesh/pycharm/provider/PythonLivingDocumentation.kt b/pycharm/src/main/kotlin/cc/unitmesh/pycharm/provider/PythonLivingDocumentation.kt index 025b7747bc..46fd42d88d 100644 --- a/pycharm/src/main/kotlin/cc/unitmesh/pycharm/provider/PythonLivingDocumentation.kt +++ b/pycharm/src/main/kotlin/cc/unitmesh/pycharm/provider/PythonLivingDocumentation.kt @@ -19,21 +19,23 @@ package cc.unitmesh.pycharm.provider import cc.unitmesh.devti.custom.document.LivingDocumentationType import cc.unitmesh.devti.provider.LivingDocumentation import com.intellij.codeInsight.daemon.impl.CollectHighlightsUtil -import com.intellij.openapi.editor.Document +import com.intellij.openapi.command.WriteCommandAction import com.intellij.openapi.editor.Editor import com.intellij.openapi.editor.SelectionModel -import com.intellij.openapi.project.Project +import com.intellij.psi.PsiDocumentManager import com.intellij.psi.PsiElement import com.intellij.psi.PsiNameIdentifierOwner import com.intellij.psi.PsiWhiteSpace import com.intellij.psi.util.PsiTreeUtil import com.intellij.util.IncorrectOperationException -import com.jetbrains.python.documentation.docstrings.PyDocstringGenerator import com.jetbrains.python.psi.* class PythonLivingDocumentation : LivingDocumentation { - override val forbiddenRules: List = listOf() + override val forbiddenRules: List = listOf( + "do not return any code, just documentation.", + "write Docstring", + ) override fun startEndString(type: LivingDocumentationType): Pair { return Pair("\"\"\"", "\"\"\"") @@ -44,8 +46,9 @@ class PythonLivingDocumentation : LivingDocumentation { throw IncorrectOperationException() } - val docstringGenerator = PyDocstringGenerator.forDocStringOwner((target as PyDocStringOwner?)!!) - docstringGenerator.buildAndInsert(newDoc, target) + WriteCommandAction.runWriteCommandAction(target.project, "Living Document", "cc.unitmesh.livingDoc", { + buildAndInsert(newDoc, target) + }); } override fun findNearestDocumentationTarget(psiElement: PsiElement): PsiNameIdentifierOwner? { @@ -113,47 +116,47 @@ class PythonLivingDocumentation : LivingDocumentation { private fun containsElement(selectionModel: SelectionModel, element: PsiElement): Boolean { return selectionModel.selectionStart <= element.textRange.startOffset && element.textRange.endOffset <= selectionModel.selectionEnd } - } -fun PyDocstringGenerator.buildAndInsert(replacementText: String, myDocStringOwner: PyDocStringOwner): PyDocStringOwner { - val project: Project = myDocStringOwner.getProject() - val elementGenerator = PyElementGenerator.getInstance(project) +fun buildAndInsert(replacementText: String, anchor: PyDocStringOwner): PyDocStringOwner { + val elementGenerator = PyElementGenerator.getInstance(anchor.project) val replacement = elementGenerator.createDocstring(replacementText) - val docStringExpression: PyStringLiteralExpression? = myDocStringOwner.getDocStringExpression() + val docStringExpression: PyStringLiteralExpression? = anchor.docStringExpression if (docStringExpression != null) { docStringExpression.replace(replacement.expression) - } else { - val container = PyUtil.`as`( - myDocStringOwner, - PyStatementListContainer::class.java - ) ?: throw IncorrectOperationException("Cannot find container for docstring, Should be a function or class") + return anchor + } - val statements = container.statementList - val indentation = PyIndentUtil.getElementIndent(statements) + val container = PyUtil.`as`( + anchor, + PyStatementListContainer::class.java + ) ?: throw IncorrectOperationException("Cannot find container for docstring, Should be a function or class") - PyUtil.updateDocumentUnblockedAndCommitted(myDocStringOwner) { document: Document -> - val beforeStatements = statements.prevSibling - var replacementWithLineBreaks = """ + val statements = container.statementList + val indentation = PyIndentUtil.getElementIndent(statements) + + val manager = PsiDocumentManager.getInstance(anchor.project) + val document = manager.getDocument(anchor.containingFile)!! + val beforeStatements = statements.prevSibling + var replacementWithLineBreaks = """ $indentation$replacementText """.trimIndent() - if (statements.statements.isNotEmpty()) { - replacementWithLineBreaks += """ + if (statements.statements.isNotEmpty()) { + replacementWithLineBreaks += """ $indentation """.trimIndent() - } - val range = beforeStatements.textRange - if (beforeStatements !is PsiWhiteSpace) { - document.insertString(range.endOffset, replacementWithLineBreaks) - } else if (statements.statements.isEmpty() && beforeStatements.textContains('\n')) { - document.insertString(range.startOffset, replacementWithLineBreaks) - } else { - document.replaceString(range.startOffset, range.endOffset, replacementWithLineBreaks) - } - } } - return myDocStringOwner + val range = beforeStatements.textRange + if (beforeStatements !is PsiWhiteSpace) { + document.insertString(range.endOffset, replacementWithLineBreaks) + } else if (statements.statements.isEmpty() && beforeStatements.textContains('\n')) { + document.insertString(range.startOffset, replacementWithLineBreaks) + } else { + document.replaceString(range.startOffset, range.endOffset, replacementWithLineBreaks) + } + + return anchor } diff --git a/src/main/kotlin/cc/unitmesh/devti/custom/document/CustomLivingDocTask.kt b/src/main/kotlin/cc/unitmesh/devti/custom/document/CustomLivingDocTask.kt index b922fb3dfe..89a042dbf5 100644 --- a/src/main/kotlin/cc/unitmesh/devti/custom/document/CustomLivingDocTask.kt +++ b/src/main/kotlin/cc/unitmesh/devti/custom/document/CustomLivingDocTask.kt @@ -27,7 +27,7 @@ class CustomLivingDocTask( val builder = CustomLivingDocPromptBuilder(editor, target, config, documentation) val prompt = builder.buildPrompt(project, target, config.prompt) - logger.warn("Prompt: $prompt") + logger.info("Prompt: $prompt") val stream = LlmFactory().create(project).stream(prompt, "") @@ -40,7 +40,7 @@ class CustomLivingDocTask( } } - logger.warn("Result: $result") + logger.info("Result: $result") documentation.updateDoc(target, result, config.type, editor) } diff --git a/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/LivingDocPromptBuilder.kt b/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/LivingDocPromptBuilder.kt index c93731793a..ae4c03f774 100644 --- a/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/LivingDocPromptBuilder.kt +++ b/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/LivingDocPromptBuilder.kt @@ -101,7 +101,7 @@ open class LivingDocPromptBuilder( } contextInstruction(llmQueryContext) - } ?: "Write documentation for given code. " + } ?: "Write documentation for given code." instruction.append(basicInstruction) @@ -116,7 +116,7 @@ open class LivingDocPromptBuilder( val startEndString = documentation.startEndString(type) instruction.append(documentation.forbiddenRules.joinToString { "\n- $it" }) - instruction.append("\n\nStart your documentation with ${startEndString.first} here, no code, and ends with `${startEndString.second}`.\n") + instruction.append("\n\nStart your documentation with ${startEndString.first} here, and ends with `${startEndString.second}`.\n") instruction.toString() } }