diff --git a/pr_agent/settings/pr_add_docs.toml b/pr_agent/settings/pr_add_docs.toml index f39eed7d..1cbe1fad 100644 --- a/pr_agent/settings/pr_add_docs.toml +++ b/pr_agent/settings/pr_add_docs.toml @@ -58,6 +58,13 @@ Code Documentation: type: integer description: |- The relevant line number from a '__new hunk__' section where the {{ docs_for_language }} should be added. + doc placement: + type: string + enum: + - before + - after + description: |- + The {{ docs_for_language }} placement relative to the relevant line (code component). documentation: type: string description: |- @@ -70,6 +77,7 @@ Code Documentation: - relevant file: |- src/file1.py relevant lines: 12 + doc placement: after documentation: |- \"\"\" This is a python docstring for func1. diff --git a/pr_agent/tools/pr_add_docs.py b/pr_agent/tools/pr_add_docs.py index 58bb1372..5859a2a7 100644 --- a/pr_agent/tools/pr_add_docs.py +++ b/pr_agent/tools/pr_add_docs.py @@ -115,8 +115,9 @@ class PRAddDocs: relevant_file = d['relevant file'].strip() relevant_line = int(d['relevant line']) # absolute position documentation = d['documentation'] + doc_placement = d['doc placement'].strip() if documentation: - new_code_snippet = self.dedent_code(relevant_file, relevant_line, documentation, + new_code_snippet = self.dedent_code(relevant_file, relevant_line, documentation, doc_placement, add_original_line=True) body = f"**Suggestion:** Proposed documentation\n```suggestion\n" + new_code_snippet + "\n```" @@ -133,7 +134,7 @@ class PRAddDocs: for code_suggestion in code_suggestions: self.git_provider.publish_code_suggestions([code_suggestion]) - def dedent_code(self, relevant_file, relevant_lines_start, new_code_snippet, add_original_line=False): + def dedent_code(self, relevant_file, relevant_lines_start, new_code_snippet, doc_placement='after', add_original_line=False): try: # dedent code snippet self.diff_files = self.git_provider.diff_files if self.git_provider.diff_files \ else self.git_provider.get_diff_files() @@ -143,14 +144,21 @@ class PRAddDocs: original_initial_line = file.head_file.splitlines()[relevant_lines_start - 1] break if original_initial_line: + if doc_placement == 'after': + line = file.head_file.splitlines()[relevant_lines_start] + else: + line = original_initial_line suggested_initial_line = new_code_snippet.splitlines()[0] - original_initial_spaces = len(original_initial_line) - len(original_initial_line.lstrip()) + original_initial_spaces = len(line) - len(line.lstrip()) suggested_initial_spaces = len(suggested_initial_line) - len(suggested_initial_line.lstrip()) delta_spaces = original_initial_spaces - suggested_initial_spaces if delta_spaces > 0: new_code_snippet = textwrap.indent(new_code_snippet, delta_spaces * " ").rstrip('\n') if add_original_line: - new_code_snippet = original_initial_line + "\n" + new_code_snippet + if doc_placement == 'after': + new_code_snippet = original_initial_line + "\n" + new_code_snippet + else: + new_code_snippet = new_code_snippet + "\n" + original_initial_line except Exception as e: if get_settings().config.verbosity_level >= 2: logging.info(f"Could not dedent code snippet for file {relevant_file}, error: {e}")