diff --git a/pr_agent/tools/pr_add_docs.py b/pr_agent/tools/pr_add_docs.py index 58bb1372..b42e9a7a 100644 --- a/pr_agent/tools/pr_add_docs.py +++ b/pr_agent/tools/pr_add_docs.py @@ -14,6 +14,22 @@ from pr_agent.git_providers.git_provider import get_main_pr_language class PRAddDocs: +""" +The PRAddDocs class is responsible for generating documentation for the code in a PR. +It uses AI models to predict the documentation for the new or edited code components. +It also handles the interaction with the git provider to get the PR details and to publish the generated documentation. + +Attributes: + git_provider: An instance of the git provider class. + main_language: The main language of the PR. + is_extended: A boolean indicating whether the extended mode is enabled. + ai_handler: An instance of the AiHandler class. + patches_diff: The diff of the PR. + prediction: The predicted documentation. + cli_mode: A boolean indicating whether the CLI mode is enabled. + vars: A dictionary containing various information about the PR. + token_handler: An instance of the TokenHandler class. +""" def __init__(self, pr_url: str, cli_mode=False, args: list = None): self.git_provider = get_git_provider()(pr_url) @@ -103,6 +119,12 @@ class PRAddDocs: return data def push_inline_docs_suggestions(self, data): + """ + This method publishes the generated documentation as inline suggestions in the PR. + + Args: + data: A dictionary containing the generated documentation. + """ code_suggestions = [] if not data['Code Documentation']: @@ -134,6 +156,18 @@ class PRAddDocs: self.git_provider.publish_code_suggestions([code_suggestion]) def dedent_code(self, relevant_file, relevant_lines_start, new_code_snippet, add_original_line=False): + """ + This method dedents the new code snippet to match the indentation of the original code. + + Args: + relevant_file: The file where the new code snippet should be added. + relevant_lines_start: The line number where the new code snippet should be added. + new_code_snippet: The new code snippet. + add_original_line: A boolean indicating whether to add the original line to the new code snippet. + + Returns: + The dedented new code snippet. + """ 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() @@ -158,6 +192,16 @@ class PRAddDocs: return new_code_snippet async def _prepare_prediction_extended(self, model: str) -> dict: + """ + This method prepares the prediction for the documentation in extended mode. + It gets the PR diff and uses an AI model to predict the documentation for each chunk of the diff. + + Args: + model: The name of the AI model to use for the prediction. + + Returns: + A dictionary containing the predicted documentation for each chunk of the diff. + """ logging.info('Getting PR diff...') patches_diff_list = get_pr_multi_diffs(self.git_provider, self.token_handler, model, max_calls=get_settings().pr_code_suggestions.max_number_of_calls)