From f1be51f5122769fc4cdd86d6d24694b37193c4d0 Mon Sep 17 00:00:00 2001
From: Bion Howard <flesheatingemu@gmail.com>
Date: Thu, 4 Jul 2024 08:23:12 -0400
Subject: [PATCH] test: multi-line f-string removal overextension; fix: same;
 fix: removal of colon

---
 tests/more_languages/group7/dataclass.py | 12 ++++++++++++
 tests/test_more_language_units.py        |  3 +++
 tree_plus_src/parse_file.py              |  8 ++------
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/tests/more_languages/group7/dataclass.py b/tests/more_languages/group7/dataclass.py
index afd2524..236fbf3 100644
--- a/tests/more_languages/group7/dataclass.py
+++ b/tests/more_languages/group7/dataclass.py
@@ -22,3 +22,15 @@ def execute(self, *args, **kwargs): ...
     output_model: DataClass
 
     def execute(self, *args, **kwargs): ...
+
+    @property
+    def edge_case(self) -> str:
+        if not self.name and not self.description:
+            return ""
+        return f"""Location: {self.config.name}
+
+{self.config.description}
+"""
+
+    def should_still_see_me(self, x: bool = True) -> "Tool":
+        pass
diff --git a/tests/test_more_language_units.py b/tests/test_more_language_units.py
index 99588f8..e2d99ac 100644
--- a/tests/test_more_language_units.py
+++ b/tests/test_more_language_units.py
@@ -2072,6 +2072,9 @@ def test_more_languages_tensorflow_flags():
     """@dataclass(frozen=True, slots=True, kw_only=True)
 class Tool(Protocol)""",
     "    def execute(self, *args, **kwargs)",
+    """    @property
+    def edge_case(self) -> str""",
+    '    def should_still_see_me(self, x: bool = True) -> "Tool"',
 ]
 
 WGSL_EXPECTATION = [
diff --git a/tree_plus_src/parse_file.py b/tree_plus_src/parse_file.py
index b343b39..8cb144e 100644
--- a/tree_plus_src/parse_file.py
+++ b/tree_plus_src/parse_file.py
@@ -1042,15 +1042,11 @@ def remove_py_comments(input_string: str) -> str:
 # from pygments.lexers.python import PythonLexer
 
 
-docstring_pattern = re.compile(r"^\s*\"\"\"[\s\S]+?\"\"\"", re.MULTILINE)
+docstring_pattern = re.compile(r"(?::)\s^\s+\"\"\"[\s\S]+?\"\"\"", re.MULTILINE)
 
 
 def remove_docstrings(source):
-    """
-    Remove docstrings from a source code string.
-    """
-    # Matches triple-quoted strings (single or double quotes)
-    return docstring_pattern.sub("", source)
+    return docstring_pattern.sub(":", source)
 
 
 def parse_py(contents: str) -> List[str]: