Add auto-fix option (#19)

README.md

@@ -107,13 +107,14 @@ int or float
 
 ## Parameters
 
-| Option             | Type          | Default   | Purpose                                                                                                     |
-| ------------------ |---------------| --------- | ----------------------------------------------------------------------------------------------------------- |
-| path               | Path          | file      | The path to the .py file or directory to analyze the functions' docstrings.                                 |
-| --ignore-dirs      | String        | "tests"   | A list of directory names to ignore while processing .py files. Separate multiple directories with a space. |
-| --error-on-warnings| Bool          | False     | If true, warnings will be treated as errors and included in the exit code count.                            |
-| --model            | String        | "gpt-4"   | The OpenAI model to use for docstring analysis. Default is 'gpt-4'.                                         |
-| --code-block-name  | String        | ""        | The name of the block you wanted audited. Leave blank to audit all code blocks.                             |
+| Option              | Type   | Default | Purpose                                                                                                     |
+|---------------------|--------|---------|-------------------------------------------------------------------------------------------------------------|
+| path                | Path   | file    | The path to the .py file or directory to analyze the functions' docstrings.                                 |
+| --ignore-dirs       | String | "tests" | A list of directory names to ignore while processing .py files. Separate multiple directories with a space. |
+| --error-on-warnings | Bool   | False   | If true, warnings will be treated as errors and included in the exit code count.                            |
+| --model             | String | "gpt-4" | The OpenAI model to use for docstring analysis. Default is 'gpt-4'.                                         |
+| --code-block-name   | String | ""      | The name of the block you wanted audited. Leave blank to audit all code blocks.                             |
+| --auto-fix          | Bool   | False   | Automatically apply the suggestions to the code. Only applied for errors, not warnings.                     |
 
 ## GitHub Action
 
@@ -129,6 +130,7 @@ To use Docstring Auditor as a GitHub Action, add the following to your workflow
       code-block-name: docstring_auditor
       model: gpt-4
       ignore-dirs: tests
+      auto-fix: false
 ```
 
 For an example of how to use Docstring Auditor, [see this workflow](https://github.com/agencyenterprise/docstring-auditor/blob/main/.github/workflows/test-workplace-action.yml).

action.yml

@@ -30,6 +30,11 @@ inputs:
     required: false
     default: "tests"
 
+  auto-fix:
+    description: "Automatically fix issues"
+    required: false
+    default: false
+
 runs:
   using: "composite"
   steps:
@@ -46,18 +51,8 @@ runs:
       shell: bash
       run: pip install --upgrade hatch
 
-    - name: pwd
-      shell: bash
-      run: |
-        pwd
-
-    - name: ls
-      shell: bash
-      run: |
-        ls
-
     - name: Execute
       shell: bash
-      run: hatch run docstring-auditor --model ${{ inputs.model }} --code-block-name ${{ inputs.code-block-name }}  --ignore-dirs ${{ inputs.ignore-dirs }} ${{ inputs.path }}
+      run: hatch run docstring-auditor --model ${{ inputs.model }} --code-block-name ${{ inputs.code-block-name }}  --ignore-dirs ${{ inputs.ignore-dirs }}  --auto-fix ${{ inputs.auto-fix }} ${{ inputs.path }}
       env:
         OPENAI_API_KEY: ${{ inputs.openaiApiKey }}

docstring_auditor/main.py

@@ -2,6 +2,7 @@
 
 import os
 import sys
+import re
 import click
 import ast
 import json
@@ -132,7 +133,7 @@ def ask_for_critique(function: str, model: str) -> Dict[str, str]:
     return response_dict
 
 
-def report_concerns(response_dict: Dict[str, str]) -> Tuple[int, int]:
+def report_concerns(response_dict: Dict[str, str]) -> Tuple[int, int, str]:
     """
     Inform the user of any concerns with the docstring.
 
@@ -175,11 +176,42 @@ def report_concerns(response_dict: Dict[str, str]) -> Tuple[int, int]:
         if solution:
             click.secho(f"A proposed solution to these concerns is:\n\n{solution}\n\n")
 
-    return error_count, warning_count
+    return error_count, warning_count, solution
+
+
+def apply_solution(file_path: str, old_function: str, new_function: str):
+    """Update the docstring of a function in a file.
+
+    This function reads the content of a file, extracts the docstrings of the old_function and new_function,
+    replaces the old docstring with the new docstring, and writes the updated content back to the file.
+
+    Parameters
+    ----------
+    file_path : str
+        The path to the file containing the function whose docstring needs to be updated.
+    old_function : str
+        The source code of the function with the old triple-quoted docstring.
+    new_function : str
+        The source code of the function with the new triple-quoted docstring.
+
+    """
+    with open(file_path, "r") as file:
+        content = file.read()
+
+    # Extract the old and new docstrings
+    old_docstring = re.search(r'""".*?"""', old_function, flags=re.DOTALL).group(0)  # type: ignore
+    new_docstring = re.search(r'""".*?"""', new_function, flags=re.DOTALL).group(0)  # type: ignore
+
+    # Replace the old docstring with the new docstring
+    updated_content = content.replace(old_docstring, new_docstring)
+
+    click.secho(f"Editing file: {file_path}", fg="red")
+    with open(file_path, "w") as file:
+        file.write(updated_content)
 
 
 def process_file(
-    file_path: str, model: str, code_block_name: str = ""
+    file_path: str, model: str, auto_fix: bool, code_block_name: str = ""
 ) -> Tuple[int, int]:
     """
     Process a single Python file and analyze its functions' and methods' docstrings.
@@ -195,6 +227,8 @@ def process_file(
         The path to the .py file to analyze the functions' and methods' docstrings.
     model : str
         The name of the OpenAI model to use for the analysis.
+    auto_fix : bool
+        Whether to automatically fix the errors and warnings found in the docstrings.
     code_block_name : str
         The name of a single block of code that you want audited, rather than all the code blocks.
         If you want all the code blocks audited, leave this blank.
@@ -211,20 +245,24 @@ def process_file(
 
     for idx, function_or_method in enumerate(functions_and_methods):
         print(
-            f"Processing function or method {idx + 1} of {len(functions_and_methods)} in file {file_path}..."
+            f"Processing code {idx + 1} of {len(functions_and_methods)} in file {file_path}..."
         )
         assert isinstance(function_or_method, str)
         critique = ask_for_critique(function_or_method, model)
-        errors, warnings = report_concerns(critique)
+        errors, warnings, solution = report_concerns(critique)
         error_count += errors
         warning_count += warnings
 
+        if auto_fix and solution:
+            apply_solution(file_path, function_or_method, solution)
+
     return error_count, warning_count
 
 
 def process_directory(
     directory_path: str,
     model: str,
+    auto_fix: bool,
     ignore_dirs: Optional[List[str]] = None,
     code_block_name: str = "",
 ) -> Tuple[int, int]:
@@ -237,6 +275,8 @@ def process_directory(
         The path to the directory containing .py files to analyze the functions' docstrings.
     model : str
         The name of the OpenAI model to use for the docstring analysis.
+    auto_fix : bool
+        Whether to automatically fix the docstring errors and warnings.
     ignore_dirs : Optional[List[str]]
         A list of directory names to ignore while processing .py files. By default, it ignores the "tests" directory.
     code_block_name : str
@@ -260,7 +300,9 @@ def process_directory(
         for file in files:
             if file.endswith(".py"):
                 file_path = os.path.join(root, file)
-                errors, warnings = process_file(file_path, model, code_block_name)
+                errors, warnings = process_file(
+                    file_path, model, auto_fix, code_block_name
+                )
                 error_count += errors
                 warning_count += warnings
 
@@ -294,12 +336,19 @@ def process_directory(
     default="",
     help="The name of a single block of code that you want audited, rather than all the code blocks.",
 )
+@click.option(
+    "--auto-fix",
+    is_flag=True,
+    default=False,
+    help="If true, the program will incorporate the suggested changes into the original file, overwriting the existing docstring.",
+)
 def docstring_auditor(
     path: str,
     ignore_dirs: List[str],
     error_on_warnings: bool,
     model: str,
     code_block_name: str,
+    auto_fix: bool,
 ):
     """
     Analyze Python functions' docstrings in a given file or directory and provide critiques and suggestions for improvement.
@@ -321,17 +370,21 @@ def docstring_auditor(
     code_block_name : str, optional
         The name of a single block of code that you want audited, rather than all the code blocks.
         If you want all the code blocks audited, leave this blank. Default is an empty string.
+    auto_fix : bool, optional
+        If true, the program will incorporate the suggested changes into the original file, overwriting the existing docstring. Default is False. Suggestions are only applied if they are associated with an error, not a warning.
 
     Returns
     -------
     None
         The function does not return any value. It prints the critiques and suggestions for the docstrings in the given file or directory.
     """
     if os.path.isfile(path):
-        error_count, warning_count = process_file(path, model, code_block_name)
+        error_count, warning_count = process_file(
+            path, model, auto_fix, code_block_name
+        )
     elif os.path.isdir(path):
         error_count, warning_count = process_directory(
-            path, model, ignore_dirs, code_block_name
+            path, model, auto_fix, ignore_dirs, code_block_name
         )
     else:
         error_text = "Invalid path. Please provide a valid file or directory path."

tests/test_autofix.py

@@ -0,0 +1,136 @@
+import pytest
+from unittest.mock import MagicMock, patch
+from click.testing import CliRunner
+from docstring_auditor.main import docstring_auditor, apply_solution
+
+
+def test_docstring_auditor_auto_fix(tmp_path):
+    # Create a temporary Python file with a sample function
+    sample_code = '''
+def sample_function(x, y):
+    """
+    Subtract two numbers together.
+    """
+    return x + y
+'''
+    temp_file = tmp_path / "sample.py"
+    temp_file.write_text(sample_code)
+
+    # Mock the ask_for_critique function to return a sample response
+    sample_response = {
+        "function": "sample_function",
+        "error": "The docstring could be improved.",
+        "warning": "The docstring could be improved.",
+        "solution": '''
+def sample_function(x, y):
+    """
+    Add two numbers together.
+
+    This function takes two integers as input and returns their sum.
+
+    Parameters
+    ----------
+    x : int
+        The first number to add.
+    y : int
+        The second number to add.
+
+    Returns
+    -------
+    int
+        The sum of x and y.
+    """
+    return x + y
+''',
+    }
+
+    with patch("docstring_auditor.main.ask_for_critique", return_value=sample_response):
+        # read temp_file
+        temp_file_contents = temp_file.read_text()
+        assert "Subtract two numbers" in temp_file_contents
+
+        # Call the docstring_auditor function with auto_fix=True
+        runner = CliRunner()
+        result = runner.invoke(docstring_auditor, ["--auto-fix", str(temp_file)])
+
+        # read temp_file
+        temp_file_contents = temp_file.read_text()
+        # assert that the words "Add two numbers" is in the contents of temp_file
+        assert "Add two numbers" in temp_file_contents
+        # assert that the words "Subtract two numbers" is not in the contents of temp_file
+        assert "Subtract two numbers" not in temp_file_contents
+
+
+def test_docstring_auditor_update_middle_function(tmp_path):
+    # Create a temporary Python file with three sample functions
+    sample_code = '''
+def function_one(x, y):
+    """
+    Add two numbers together.
+    """
+    return x + y
+
+def function_two(x, y):
+    """
+    Subtract two numbers together.
+    """
+    return x - y
+
+def function_three(x, y):
+    """
+    Multiply two numbers together.
+    """
+    return x * y
+'''
+    temp_file = tmp_path / "sample.py"
+    temp_file.write_text(sample_code)
+
+    # Mock the ask_for_critique function to return a sample response for function_two
+    sample_response = {
+        "function": "function_two",
+        "error": "The docstring could be improved.",
+        "warning": "The docstring could be improved.",
+        "solution": '''
+def function_two(x, y):
+    """
+    Subtract two numbers.
+
+    This function takes two integers as input and returns the result of subtracting the second number from the first.
+
+    Parameters
+    ----------
+    x : int
+        The number to subtract from.
+    y : int
+        The number to subtract.
+
+    Returns
+    -------
+    int
+        The result of x - y.
+    """
+    return x - y
+''',
+    }
+
+    with patch("docstring_auditor.main.ask_for_critique", return_value=sample_response):
+        # read temp_file
+        temp_file_contents = temp_file.read_text()
+        assert "Subtract two numbers together." in temp_file_contents
+
+        # Call the docstring_auditor function with auto_fix=True and code_block_name="function_two"
+        runner = CliRunner()
+        result = runner.invoke(
+            docstring_auditor,
+            ["--auto-fix", "--code-block-name", "function_two", str(temp_file)],
+        )
+
+        # read temp_file
+        temp_file_contents = temp_file.read_text()
+        # assert that the words "Subtract two numbers." is in the contents of temp_file
+        assert "Subtract two numbers." in temp_file_contents
+        # assert that the words "Subtract two numbers together." is not in the contents of temp_file
+        assert "Subtract two numbers together." not in temp_file_contents
+        # assert that the other functions' docstrings remain unchanged
+        assert "Add two numbers together." in temp_file_contents
+        assert "Multiply two numbers together." in temp_file_contents

tests/test_reporter.py

@@ -18,7 +18,7 @@ def mock_secho(*args, **kwargs):
         "solution": "",
     }
     result = report_concerns(response_dict)
-    assert result == (0, 0)
+    assert result == (0, 0, "")
     assert len(captured_output) == 1
     assert "No concerns found" in captured_output[0]
 
@@ -38,7 +38,7 @@ def mock_secho(*args, **kwargs):
         "solution": "Updated docstring.",
     }
     result = report_concerns(response_dict)
-    assert result == (1, 0)
+    assert result == (1, 0, "Updated docstring.")
     assert len(captured_output) == 3
     assert "An error was found" in captured_output[0]
     assert "An error occurred." in captured_output[1]
@@ -60,7 +60,7 @@ def mock_secho(*args, **kwargs):
         "solution": "",
     }
     result = report_concerns(response_dict)
-    assert result == (0, 1)
+    assert result == (0, 1, "")
     assert len(captured_output) == 2
     assert "A warning was found" in captured_output[0]
     assert "A warning occurred." in captured_output[1]
@@ -82,7 +82,7 @@ def mock_secho(*args, **kwargs):
     }
     result = report_concerns(response_dict)
     print(captured_output)
-    assert result == (1, 1)
+    assert result == (1, 1, "Updated docstring.")
     assert len(captured_output) == 5
     assert "An error was found" in captured_output[0]
     assert "An error occurred." in captured_output[1]