diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py index e5d13475..58f0ccb9 100644 --- a/pr_agent/algo/utils.py +++ b/pr_agent/algo/utils.py @@ -297,6 +297,21 @@ def load_yaml(review_text: str) -> dict: def try_fix_yaml(review_text: str) -> dict: review_text_lines = review_text.split('\n') + + # first fallback - try to convert 'relevant line: ...' to relevant line: |-\n ...' + review_text_lines_copy = review_text_lines.copy() + for i in range(0, len(review_text_lines_copy)): + if 'relevant line:' in review_text_lines_copy[i] and not '|-' in review_text_lines_copy[i]: + review_text_lines_copy[i] = review_text_lines_copy[i].replace('relevant line: ', + 'relevant line: |-\n ') + try: + data = yaml.load('\n'.join(review_text_lines_copy), Loader=yaml.SafeLoader) + get_logger().info(f"Successfully parsed AI prediction after adding |-\n to relevant line") + return data + except: + get_logger().debug(f"Failed to parse AI prediction after adding |-\n to relevant line") + + # second fallback - try to remove last lines data = {} for i in range(1, len(review_text_lines)): review_text_lines_tmp = '\n'.join(review_text_lines[:-i]) diff --git a/tests/unittest/try_fix_yaml.py b/tests/unittest/try_fix_yaml.py new file mode 100644 index 00000000..fe37bb1c --- /dev/null +++ b/tests/unittest/try_fix_yaml.py @@ -0,0 +1,31 @@ + +# Generated by CodiumAI +from pr_agent.algo.utils import try_fix_yaml + + +import pytest + +class TestTryFixYaml: + + # The function successfully parses a valid YAML string. + def test_valid_yaml(self): + review_text = "key: value\n" + expected_output = {"key": "value"} + assert try_fix_yaml(review_text) == expected_output + + # The function adds '|-' to 'relevant line:' if it is not already present and successfully parses the YAML string. + def test_add_relevant_line(self): + review_text = "relevant line: value: 3\n" + expected_output = {"relevant line": "value: 3"} + assert try_fix_yaml(review_text) == expected_output + + # The function removes the last line(s) of the YAML string and successfully parses the YAML string. + def test_remove_last_line(self): + review_text = "key: value\nextra invalid line\n" + expected_output = {"key": "value"} + assert try_fix_yaml(review_text) == expected_output + + # The YAML string is empty. + def test_empty_yaml_fixed(self): + review_text = "" + assert try_fix_yaml(review_text) is None \ No newline at end of file