PR: fix: docstring duplication

Artemonim · web-flow · commit 8b7fa17259b6 · 2025-07-06T12:25:55.000+03:00
### Fixed - **Python Docstring Cleaning**: (fixes #11) - **C-Style Comment Handling**:
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,7 +16,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [NextRelease]
 
--   **Something great**: for sure
+### Fixed
+
+-   **Python Docstring Cleaning**: Improved the `remove_agent_docstring` function to better handle Python docstrings by preserving manual content while removing auto-generated table of contents. The function now correctly identifies and removes only the auto-generated content while maintaining the structure of existing manual docstrings. (fixes #11)
+-   **C-Style Comment Handling**: Enhanced the docstring removal logic for C-style languages (Kotlin, Java, Go, etc.) to be more flexible with comment formatting variations, ensuring proper detection and removal of auto-generated content across different comment styles.
 
 ## [1.3.2]
 
diff --git a/agent_docstrings/languages/common.py b/agent_docstrings/languages/common.py
@@ -1,16 +1,15 @@
 """
     --- AUTO-GENERATED DOCSTRING ---
-    Table of content is automatically generated by Agent Docstrings v1.3.0
+    Table of content is automatically generated by Agent Docstrings v1.3.2
     
     Classes/Functions:
-        - SignatureInfo (line 20):
-        - ClassInfo (line 26):
-        - CommentStyle (line 34):
-        - remove_agent_docstring(text: str, language: str) -> str (line 57)
+        - SignatureInfo (line 17):
+        - ClassInfo (line 21):
+        - CommentStyle (line 27):
+        - remove_agent_docstring(text: str, language: str) -> str (line 46)
     --- END AUTO-GENERATED DOCSTRING ---
 """
 from __future__ import annotations
-
 import re
 from typing import List, Tuple, Dict, NamedTuple
 
@@ -22,23 +21,20 @@ class SignatureInfo(NamedTuple):
     signature: str
     line: int
 
-
 class ClassInfo(NamedTuple):
     """Stores information about a parsed class, including its methods."""
     name: str
     line: int
     methods: List[SignatureInfo]
     inner_classes: List["ClassInfo"]
 
-
 class CommentStyle(NamedTuple):
     """Stores language-specific comment formatting information."""
     start: str
     end: str
     prefix: str
     indent: str
 
-
 COMMENT_STYLES: Dict[str, CommentStyle] = {
     "python": CommentStyle('"""', '"""', "    ", "    "),
     "kotlin": CommentStyle('/**', ' */', ' * ', "    "),
@@ -53,63 +49,43 @@ class CommentStyle(NamedTuple):
     "delphi": CommentStyle('(*', '*)', ' * ', "  "),
 }
 
-
 def remove_agent_docstring(text: str, language: str) -> str:
-    """Remove a previously generated docstring from *text*.
-
-    The search uses language-specific comment patterns to find a block
-    containing DOCSTRING_START_MARKER and DOCSTRING_END_MARKER at the
-    beginning of the file, and removes it.
-
-    Args:
-        text (str): Full contents of the source file.
-        language (str): Canonical language name (e.g. ``"python"``) used
-            to pick the correct comment delimiters from
-            :data:`COMMENT_STYLES`.
-
-    Returns:
-        str: *text* without the agent docstring block. If no such
-        docstring is detected, *text* is returned unchanged.
-    """
+    """Remove a previously generated docstring from *text*."""
     style = COMMENT_STYLES[language]
-    
-    # ! Create a more flexible pattern that can match various formats
     start_marker_escaped = re.escape(DOCSTRING_START_MARKER)
     end_marker_escaped = re.escape(DOCSTRING_END_MARKER)
-    
     if language == "python":
-        # * Python uses triple quotes - check for new format first
-        pattern = re.compile(
-            rf'^\s*"""\s*{start_marker_escaped}.*?{end_marker_escaped}\s*"""\s*\n?',
-            re.DOTALL
-        )
-        match = pattern.search(text)
-        if match:
-            return text[match.end():]
-        
-        # * Also check for old format (without proper markers)
-        old_format_pattern = re.compile(
-            rf'^\s*"""\s*Classes/Functions:.*?"""\s*\n?',
-            re.DOTALL
-        )
-        match = old_format_pattern.search(text)
-        if match:
-            return text[match.end():]
+        def replacer(match):
+            docstring_content = match.group(0)
+            auto_content_pattern = re.compile(
+                rf'\s*{start_marker_escaped}[\s\S]*?{end_marker_escaped}\s*?\n?',
+                re.DOTALL
+            )
+            cleaned_docstring = auto_content_pattern.sub('', docstring_content)
+            temp_cleaned = cleaned_docstring.replace('"""', '').replace("'''", '').strip()
+            if not temp_cleaned:
+                return ''  # Remove empty docstring
+            # Ensure single newline padding for non-empty manual comments
+            return f'"""\n{temp_cleaned}\n"""'
+        docstring_pattern = re.compile(r'^\s*("""[\s\S]*?"""|'r"'''[\s\S]*?''')")
+        # Iteratively clean the text
+        cleaned_text = docstring_pattern.sub(replacer, text)
+        cleaned_text = docstring_pattern.sub(replacer, cleaned_text) # Run again to handle adjacent blocks
+        # Collapse whitespace and return
+        return cleaned_text.strip()
     else:
-        # * For C-style comments, be more flexible with the format
-        # * Handle both compact (/**---...---*/) and expanded formats
+        # For C-style comments, be more flexible with the format
+        # Handle both compact (/**---...---*/) and expanded formats
         start_escaped = re.escape(style.start.rstrip())  # Remove trailing spaces
-        
-        # * Handle different possible endings (with or without space before *)
+        # Handle different possible endings (with or without space before *)
         end_patterns = [
             re.escape(style.end),  # Original format with space
             re.escape(style.end.strip()),  # Without space
         ]
-        
-        # * Try each possible end pattern
+        # Try each possible end pattern
         for end_pattern in end_patterns:
             pattern = re.compile(
-                rf'^\s*{start_escaped}.*?{start_marker_escaped}.*?{end_marker_escaped}.*?{end_pattern}\s*\n?',
+                rf'^\s*{start_escaped}[\s\S]*?{start_marker_escaped}[\s\S]*?{end_marker_escaped}[\s\S]*?{end_pattern}\s*\n?',
                 re.DOTALL
             )
             match = pattern.search(text)
diff --git a/tests/test_common.py b/tests/test_common.py
@@ -2,34 +2,31 @@
 
 """
     --- AUTO-GENERATED DOCSTRING ---
-    Table of content is automatically generated by Agent Docstrings v1.3.1
+    Table of content is automatically generated by Agent Docstrings v1.3.2
     
     Classes/Functions:
-        - TestDataClasses (line 40):
-            - test_signature_info_creation() -> None (line 43)
-            - test_class_info_creation() -> None (line 49)
-            - test_comment_style_creation() -> None (line 68)
-        - TestCommentStyles (line 77):
-            - test_all_supported_languages_have_styles() -> None (line 80)
-            - test_comment_style_values(language: str, expected_start: str, expected_end: str, expected_prefix: str, expected_indent: str) -> None (line 97)
-        - TestHeaderStripping (line 113):
-            - test_strip_python_header() -> None (line 116)
-            - test_strip_block_comment_header() -> None (line 141)
-            - test_strip_c_style_comment_header() -> None (line 159)
-            - test_no_header_to_strip() -> None (line 177)
-            - test_preserve_shebang_when_stripping() -> None (line 186)
-            - test_strip_header_with_various_whitespace() -> None (line 199)
-            - test_strip_only_first_matching_header() -> None (line 207)
-            - test_strip_header_edge_cases() -> None (line 223)
-            - test_header_not_at_start() -> None (line 235)
-            - test_invalid_language_patterns(language: str) -> None (line 249)
+        - TestDataClasses (line 37):
+            - test_signature_info_creation() -> None (line 39)
+            - test_class_info_creation() -> None (line 44)
+            - test_comment_style_creation() -> None (line 60)
+        - TestCommentStyles (line 67):
+            - test_all_supported_languages_have_styles() -> None (line 69)
+            - test_comment_style_values(language: str, expected_start: str, expected_end: str, expected_prefix: str, expected_indent: str) -> None (line 85)
+        - TestHeaderStripping (line 99):
+            - test_strip_python_header() -> None (line 101)
+            - test_strip_block_comment_header() -> None (line 121)
+            - test_strip_c_style_comment_header() -> None (line 136)
+            - test_no_header_to_strip() -> None (line 151)
+            - test_preserve_shebang_when_stripping() -> None (line 158)
+            - test_strip_header_with_various_whitespace() -> None (line 169)
+            - test_strip_only_first_matching_header() -> None (line 175)
+            - test_strip_header_edge_cases() -> None (line 189)
+            - test_header_not_at_start() -> None (line 198)
+            - test_invalid_language_patterns(language: str) -> None (line 209)
     --- END AUTO-GENERATED DOCSTRING ---
 Tests for agent_docstrings.languages.common module.
 """
-
-
 import pytest
-
 from agent_docstrings.languages.common import (
     COMMENT_STYLES,
     ClassInfo,
@@ -40,7 +37,6 @@
     DOCSTRING_END_MARKER,
 )
 
-
 class TestDataClasses:
     """Tests for data classes used in parsing."""
 
@@ -77,7 +73,6 @@ def test_comment_style_creation(self) -> None:
         assert style.prefix == " * "
         assert style.indent == "  "
 
-
 class TestCommentStyles:
     """Tests for comment style definitions."""
 
@@ -113,7 +108,6 @@ def test_comment_style_values(
         assert style.prefix == expected_prefix
         assert style.indent == expected_indent
 
-
 class TestHeaderStripping:
     """Tests for remove_agent_docstring function."""
 
diff --git a/tests/test_docstring_duplication.py b/tests/test_docstring_duplication.py
@@ -0,0 +1,170 @@
+"""
+    --- AUTO-GENERATED DOCSTRING ---
+    Table of content is automatically generated by Agent Docstrings v1.3.2
+    
+    Classes/Functions:
+        - test_no_docstring_duplication_on_repeated_runs(source_processor) -> None (line 15)
+        - test_manual_docstring_preservation_with_auto_generation(source_processor) -> None (line 56)
+        - test_existing_auto_docstring_replacement(source_processor) -> None (line 100)
+        - test_multiple_auto_docstring_removal(source_processor) -> None (line 135)
+    --- END AUTO-GENERATED DOCSTRING ---
+"""
+import pytest
+import re
+from textwrap import dedent
+
+def test_no_docstring_duplication_on_repeated_runs(source_processor) -> None:
+    """
+    Test that running the docstring generator multiple times on the same file
+    does not create duplicate auto-generated docstrings.
+    This test simulates the scenario where a file with manual docstring
+    gets processed multiple times, ensuring no double docstrings are created.
+    """
+    # * Initial file with manual docstring
+    initial_content = dedent('''
+        """
+        Human comments
+        This is a manual docstring that should be preserved.
+        """
+        def test_function():
+            """This is a function docstring."""
+            return "test"
+        class TestClass:
+            def method(self):
+                return "method"
+    ''').strip()
+    # * First run - should generate auto docstring and merge with manual
+    result_content_1, lines_1, _ = source_processor("test_duplication.py", initial_content)
+    # * Verify that auto-generated docstring was added
+    assert "--- AUTO-GENERATED DOCSTRING ---" in result_content_1
+    assert "Human comments" in result_content_1  # Manual content preserved
+    assert "test_function()" in result_content_1  # Auto-generated content added
+    # * Count auto-generated docstring markers
+    auto_markers_1 = result_content_1.count("--- AUTO-GENERATED DOCSTRING ---")
+    assert auto_markers_1 == 1, f"Expected 1 auto docstring marker, found {auto_markers_1}"
+    # * Second run - should not create duplicate auto docstrings
+    result_content_2, lines_2, _ = source_processor("test_duplication.py", result_content_1)
+    # * Verify no duplication occurred
+    auto_markers_2 = result_content_2.count("--- AUTO-GENERATED DOCSTRING ---")
+    assert auto_markers_2 == 1, f"Expected 1 auto docstring marker after second run, found {auto_markers_2}"
+    # * Verify manual content is still preserved
+    assert "Human comments" in result_content_2
+    assert "This is a manual docstring that should be preserved." in result_content_2
+    # * Verify auto-generated content is still present
+    assert "test_function()" in result_content_2
+    assert "TestClass" in result_content_2
+    assert "method()" in result_content_2
+
+def test_manual_docstring_preservation_with_auto_generation(source_processor) -> None:
+    """
+    Test that manual docstrings are properly preserved when auto-generating
+    docstrings, and that the structure is correct.
+    """
+    # * File with manual docstring only
+    initial_content = dedent('''
+        """
+        This is a manual module docstring.
+        It should be preserved and merged with auto-generated content.
+        """
+        def function_one():
+            pass
+        def function_two():
+            pass
+    ''').strip()
+    result_content, lines, _ = source_processor("test_manual_preservation.py", initial_content)
+    # * Verify structure: manual content should come after auto-generated content
+    lines_list = result_content.split('\n')
+    # * Find the docstring boundaries
+    docstring_start = None
+    docstring_end = None
+    manual_content_found = False
+    for i, line in enumerate(lines_list):
+        if line.strip() == '"""' and docstring_start is None:
+            docstring_start = i
+        elif line.strip() == '"""' and docstring_start is not None:
+            docstring_end = i
+            break
+    assert docstring_start is not None, "Docstring start not found"
+    assert docstring_end is not None, "Docstring end not found"
+    # * Extract docstring content
+    docstring_content = lines_list[docstring_start:docstring_end + 1]
+    docstring_text = '\n'.join(docstring_content)
+    # * Verify auto-generated content is first
+    assert "--- AUTO-GENERATED DOCSTRING ---" in docstring_text
+    assert "function_one()" in docstring_text
+    assert "function_two()" in docstring_text
+    # * Verify manual content is preserved
+    assert "This is a manual module docstring." in docstring_text
+    assert "It should be preserved and merged with auto-generated content." in docstring_text
+    # * Verify only one docstring block exists
+    docstring_blocks = result_content.count('"""')
+    assert docstring_blocks == 2, f"Expected 2 triple quotes (start and end), found {docstring_blocks}"
+
+def test_existing_auto_docstring_replacement(source_processor) -> None:
+    """
+    Test that existing auto-generated docstrings are properly replaced
+    when the file is processed again.
+    """
+    # * File with existing auto-generated docstring
+    initial_content = dedent('''
+        """
+        --- AUTO-GENERATED DOCSTRING ---
+        Table of content is automatically generated by Agent Docstrings v1.3.1
+        Classes/Functions:
+            - old_function() (line 8)
+        --- END AUTO-GENERATED DOCSTRING ---
+        """
+        def old_function():
+            pass
+        def new_function():
+            pass
+    ''').strip()
+    result_content, lines, _ = source_processor("test_replacement.py", initial_content)
+    # * Find the docstring in the result
+    docstring_match = re.search(r'"""[\s\S]*?"""', result_content)
+    assert docstring_match, "Could not find docstring in processed file"
+    docstring_text = docstring_match.group(0)
+    # * Verify new content is in the docstring
+    assert "old_function()" in docstring_text
+    assert "new_function()" in docstring_text
+    # * Verify only one auto-generated docstring exists in the whole file
+    auto_markers = result_content.count("--- AUTO-GENERATED DOCSTRING ---")
+    assert auto_markers == 1, f"Expected 1 auto docstring marker, found {auto_markers}"
+    # * Verify the version is updated in the docstring
+    assert "Agent Docstrings v1.3.2" in docstring_text
+    # * Verify that old_function is mentioned only once *within the docstring*
+    assert docstring_text.count("old_function()") == 1, "Function should appear only once in docstring"
+    assert docstring_text.count("new_function()") == 1, "Function should appear only once in docstring"
+
+def test_multiple_auto_docstring_removal(source_processor) -> None:
+    """
+    Test that multiple auto-generated docstrings are properly removed
+    and replaced with a single one.
+    """
+    # * File with multiple auto-generated docstrings (simulating a bug)
+    initial_content = dedent('''
+        """
+        --- AUTO-GENERATED DOCSTRING ---
+        Table of content is automatically generated by Agent Docstrings v1.3.1
+        --- END AUTO-GENERATED DOCSTRING ---
+        """
+        """
+        --- AUTO-GENERATED DOCSTRING ---
+        Table of content is automatically generated by Agent Docstrings v1.3.2
+        --- END AUTO-GENERATED DOCSTRING ---
+        Human comments
+        """
+        def test_function():
+            return "test"
+    ''').strip()
+    result_content, lines, _ = source_processor("test_multiple_removal.py", initial_content)
+    # * Verify only one auto-generated docstring exists
+    auto_markers = result_content.count("--- AUTO-GENERATED DOCSTRING ---")
+    assert auto_markers == 1, f"Expected 1 auto docstring marker, found {auto_markers}"
+    # * Verify manual content is preserved
+    assert "Human comments" in result_content
+    # * Verify function is documented
+    assert "test_function()" in result_content
+    # * Verify that there is only one docstring block in the final output
+    docstring_blocks = re.findall(r'"""[\s\S]*?"""', result_content)
+    assert len(docstring_blocks) == 1, f"Expected 1 docstring block, found {len(docstring_blocks)}" 
diff --git a/tests/test_whitespace_preservation.py b/tests/test_whitespace_preservation.py
