xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-04 21:00:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #464 from Codium-ai/tr/more_protections

Browse Source 
Refactor YAML parsing for improved error handling
This commit is contained in:
mrT23
2023-11-20 02:28:57 -08:00
committed by GitHub
parent 26cc26129c 3f76d95495
commit b3129c7dd9
2 changed files with 19 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
pr_agent/algo/utils.py
View File

@ -304,14 +304,14 @@ def try_fix_yaml(response_text: str) -> dict:
for i in range(0, len(response_text_lines_copy)): for i in range(0, len(response_text_lines_copy)):
for key in keys: for key in keys:
if key in response_text_lines_copy[i] and not '|-' in response_text_lines_copy[i]: if key in response_text_lines_copy[i] and not '|-' in response_text_lines_copy[i]:
response_text_lines_copy[i] = response_text_lines_copy[i].replace(f'{key}: ', response_text_lines_copy[i] = response_text_lines_copy[i].replace(f'{key}',
f'{key}: |-\n ') f'{key} |-\n ')
try: try:
data = yaml.safe_load('\n'.join(response_text_lines_copy)) data = yaml.safe_load('\n'.join(response_text_lines_copy))
get_logger().info(f"Successfully parsed AI prediction after adding |-\n") get_logger().info(f"Successfully parsed AI prediction after adding |-\n")
return data return data
except: except:
get_logger().debug(f"Failed to parse AI prediction after adding |-\n") get_logger().info(f"Failed to parse AI prediction after adding |-\n")
# second fallback - try to remove last lines # second fallback - try to remove last lines
data = {} data = {}

21
tests/unittest/test_load_yaml.py
View File

@ -2,6 +2,9 @@
# Generated by CodiumAI # Generated by CodiumAI
import pytest import pytest
import yaml
from yaml.scanner import ScannerError
from pr_agent.algo.utils import load_yaml from pr_agent.algo.utils import load_yaml
@ -26,15 +29,23 @@ PR Feedback:
Code feedback: Code feedback:
- relevant file: pr_agent/settings/pr_description_prompts.toml - relevant file: pr_agent/settings/pr_description_prompts.toml
suggestion: Consider using a more descriptive variable name than 'user' for the command prompt. A more descriptive name would make the code more readable and maintainable. [medium] suggestion: Consider using a more descriptive variable name than 'user' for the command prompt. A more descriptive name would make the code more readable and maintainable. [medium]
relevant line: 'user="""PR Info:' relevant line: user="""PR Info: aaa
Security concerns: No''' Security concerns: No'''
expected_output = {'PR Analysis': {'Main theme': 'Enhancing the `/describe` command prompt by adding title and description', 'Type of PR': 'Enhancement', 'Relevant tests added': False, 'Focused PR': 'Yes, the PR is focused on enhancing the `/describe` command prompt.'}, 'PR Feedback': {'General suggestions': 'The PR seems to be well-structured and focused on a specific enhancement. However, it would be beneficial to add tests to ensure the new feature works as expected.', 'Code feedback': [{'relevant file': 'pr_agent/settings/pr_description_prompts.toml', 'suggestion': "Consider using a more descriptive variable name than 'user' for the command prompt. A more descriptive name would make the code more readable and maintainable. [medium]", 'relevant line': 'user="""PR Info:'}], 'Security concerns': False}} with pytest.raises(ScannerError):
yaml.safe_load(yaml_str)
expected_output = {'PR Analysis': {'Main theme': 'Enhancing the `/describe` command prompt by adding title and description', 'Type of PR': 'Enhancement', 'Relevant tests added': False, 'Focused PR': 'Yes, the PR is focused on enhancing the `/describe` command prompt.'}, 'PR Feedback': {'General suggestions': 'The PR seems to be well-structured and focused on a specific enhancement. However, it would be beneficial to add tests to ensure the new feature works as expected.', 'Code feedback': [{'relevant file': 'pr_agent/settings/pr_description_prompts.toml', 'suggestion': "Consider using a more descriptive variable name than 'user' for the command prompt. A more descriptive name would make the code more readable and maintainable. [medium]", 'relevant line': 'user="""PR Info: aaa'}], 'Security concerns': False}}
assert load_yaml(yaml_str) == expected_output assert load_yaml(yaml_str) == expected_output
def test_load_invalid_yaml2(self): def test_load_invalid_yaml2(self):
yaml_str = '''\ yaml_str = '''\
- relevant file: src/app.py - relevant file: src/app.py:
suggestion content: The print statement is outside inside the if __name__ == \ suggestion content: The print statement is outside inside the if __name__ ==: \
''' '''
expected_output =[{'relevant file': 'src/app.py', 'suggestion content': 'The print statement is outside inside the if __name__ =='}] with pytest.raises(ScannerError):
yaml.safe_load(yaml_str)
expected_output =[{'relevant file': 'src/app.py:',
'suggestion content': 'The print statement is outside inside the if __name__ ==: '}]
assert load_yaml(yaml_str) == expected_output assert load_yaml(yaml_str) == expected_output