155 lines
4.2 KiB
Python
155 lines
4.2 KiB
Python
import os
|
||
import re
|
||
import sys
|
||
|
||
from devchat.llm import chat
|
||
|
||
from lib.ide_service import IDEService
|
||
|
||
|
||
def get_selected_code():
|
||
"""
|
||
Retrieves the selected lines of code from the user's selection.
|
||
|
||
This function extracts the text selected by the user in their IDE or text editor.
|
||
If no text has been selected, it prints an error message to stderr and exits the
|
||
program with a non-zero status indicating failure.
|
||
|
||
Returns:
|
||
dict: A dictionary containing the key 'selectedText' with the selected text
|
||
as its value. If no text is selected, the program exits.
|
||
"""
|
||
selected_data = IDEService().get_selected_range().dict()
|
||
|
||
miss_selected_error = "Please select some text."
|
||
if selected_data["range"]["start"] == selected_data["range"]["end"]:
|
||
readme_path = os.path.join(os.path.dirname(__file__), "README.md")
|
||
if os.path.exists(readme_path):
|
||
with open(readme_path, "r", encoding="utf-8") as f:
|
||
readme_text = f.read()
|
||
print(readme_text)
|
||
sys.exit(0)
|
||
|
||
print(miss_selected_error, file=sys.stderr, flush=True)
|
||
sys.exit(-1)
|
||
|
||
return selected_data
|
||
|
||
|
||
def get_visible_code():
|
||
"""
|
||
Retrieves visible code from the visible_lines function.
|
||
|
||
Returns:
|
||
visible_data: The visible code retrieved from the visible_lines function.
|
||
"""
|
||
visible_data = IDEService().get_visible_range().dict()
|
||
return visible_data
|
||
|
||
|
||
PROMPT = prompt = """
|
||
Your task is:
|
||
|
||
Try to find potential bugs in the selected code and fix it. \
|
||
Please only fix and modify the selected portion of the code and giving out the fixed code. \
|
||
Do not show the original code in your results. \
|
||
Explain the bugs and fixes at last.
|
||
|
||
Following the task requirements, please ensure that the revised code segment \
|
||
maintains the same indentation as the selected code to seamlessly integrate with \
|
||
the existing code structure and maintain correct syntax.
|
||
Giving the results in the following format:
|
||
|
||
Fixed code:
|
||
```language
|
||
The fixed code goes here...
|
||
```
|
||
|
||
Bugs and fixes:
|
||
1. Bug1 and its fixes...;
|
||
2. Bug2 and its fixes...;
|
||
...
|
||
|
||
Here is the relevant context information for your reference:
|
||
1. Selected portion of the code: {selected_text}
|
||
"""
|
||
|
||
|
||
PROMPT_ZH = prompt = """
|
||
你的任务是：
|
||
|
||
尝试找到被选中代码中的潜在错误并修复它。 \
|
||
请仅修复和修改被选中的代码并给出修复后的代码。 \
|
||
可见代码仅作为你理解的 context，不要修改也不要输出可见代码 \
|
||
不要在结果中显示原始代码。 \
|
||
最后解释一下错误和修复。
|
||
|
||
按照任务要求，请确保修改后的代码段保持与被选中代码相同的缩进，以与现有的代码结构无缝集成并保持正确的语法。
|
||
以以下格式给出结果：
|
||
|
||
已修复代码：
|
||
```language
|
||
已修复代码放在这里...
|
||
```
|
||
|
||
问题及修复：
|
||
|
||
1. Bug1 及其修复...；
|
||
2. Bug2 及其修复...；
|
||
...
|
||
|
||
以下是相关上下文信息，供您参考：
|
||
1. 编辑器中被选中的代码：{selected_text}
|
||
"""
|
||
|
||
|
||
def get_prompt():
|
||
ide_language = IDEService().ide_language()
|
||
return PROMPT_ZH if ide_language == "zh" else PROMPT
|
||
|
||
|
||
@chat(prompt=get_prompt(), stream_out=True)
|
||
# pylint: disable=unused-argument
|
||
def fix_bugs(selected_text, visible_text):
|
||
"""
|
||
call ai to rewrite selected code
|
||
"""
|
||
pass # pylint: disable=unnecessary-pass
|
||
|
||
|
||
def extract_markdown_block(text):
|
||
"""
|
||
Extracts the first Markdown code block from the given text without the language specifier.
|
||
|
||
:param text: A string containing Markdown text
|
||
:return: The content of the first Markdown code block, or None if not found
|
||
"""
|
||
pattern = r"```(?:\w+)?\s*\n(.*?)\n```"
|
||
match = re.search(pattern, text, re.DOTALL)
|
||
|
||
if match:
|
||
block_content = match.group(1)
|
||
return block_content
|
||
else:
|
||
return None
|
||
|
||
|
||
def main():
|
||
# prepare code
|
||
selected_text = get_selected_code()
|
||
visible_text = get_visible_code()
|
||
|
||
# rewrite
|
||
response = fix_bugs(selected_text=selected_text, visible_text=visible_text)
|
||
|
||
# apply new code to editor
|
||
new_code = extract_markdown_block(response)
|
||
if new_code is not None:
|
||
IDEService().diff_apply("", new_code)
|
||
|
||
sys.exit(0)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|