Handle docstrings in single quoted strings correctly

pydocstringformatter/formatting/base.py

@@ -1,5 +1,7 @@
 import abc
+import re
 import tokenize
+from typing import Literal
 
 
 class Formatter:
@@ -47,3 +49,42 @@ def treat_token(self, tokeninfo: tokenize.TokenInfo) -> tokenize.TokenInfo:
             tokeninfo.end,
             tokeninfo.line,
         )
+
+
+class StringAndQuotesFormatter(Formatter):
+    """Base class for string formatter that needs access to the quotes."""
+
+    quotes_regex = re.compile(r"""['"]{1,3}""")
+    """Pattern to match against opening quotes."""
+
+    @abc.abstractmethod
+    def _treat_string(
+        self,
+        tokeninfo: tokenize.TokenInfo,
+        indent_length: int,
+        quotes: str,
+        quotes_length: Literal[1, 3],
+    ) -> str:
+        """Return a modified string."""
+
+    def treat_token(self, tokeninfo: tokenize.TokenInfo) -> tokenize.TokenInfo:
+        # Get the quotes used for this docstring
+        match = re.match(self.quotes_regex, tokeninfo.string)
+        assert match
+        quotes = match.group()
+
+        quotes_length = len(quotes)
+        assert quotes_length in {1, 3}
+
+        return tokenize.TokenInfo(
+            tokeninfo.type,
+            self._treat_string(
+                tokeninfo,
+                tokeninfo.start[1],
+                quotes,
+                quotes_length,  # type: ignore[arg-type]
+            ),
+            tokeninfo.start,
+            tokeninfo.end,
+            tokeninfo.line,
+        )

pydocstringformatter/formatting/formatter.py

@@ -1,7 +1,11 @@
 import re
 import tokenize
+from typing import Literal
 
-from pydocstringformatter.formatting.base import StringFormatter
+from pydocstringformatter.formatting.base import (
+    StringAndQuotesFormatter,
+    StringFormatter,
+)
 
 
 class BeginningQuotesFormatter(StringFormatter):
@@ -20,7 +24,7 @@ class CapitalizeFirstLetterFormatter(StringFormatter):
     """Capitalize the first letter of the docstring if appropriate."""
 
     name = "capitalize-first-letter"
-    first_letter_re = re.compile(r"""['"]{3}\s*(\w)""")
+    first_letter_re = re.compile(r"""['"]{1,3}\s*(\w)""")
 
     def _treat_string(self, tokeninfo: tokenize.TokenInfo, _: int) -> str:
         new_string = None
@@ -57,17 +61,23 @@ def _treat_string(self, tokeninfo: tokenize.TokenInfo, _: int) -> str:
         return new_string
 
 
-class FinalPeriodFormatter(StringFormatter):
+class FinalPeriodFormatter(StringAndQuotesFormatter):
     """Add a period to the end of single line docstrings and summaries."""
 
     name = "final-period"
 
-    def _treat_string(self, tokeninfo: tokenize.TokenInfo, _: int) -> str:
+    def _treat_string(
+        self,
+        tokeninfo: tokenize.TokenInfo,
+        _: int,
+        quotes: str,
+        quotes_length: Literal[1, 3],
+    ) -> str:
         """Add a period to the end of single-line docstrings and summaries."""
         # Handle single line docstrings
         if not tokeninfo.string.count("\n"):
-            if tokeninfo.string[-4] != ".":
-                return tokeninfo.string[:-3] + "." + tokeninfo.string[-3:]
+            if tokeninfo.string[-quotes_length - 1] != ".":
+                return tokeninfo.string[:-quotes_length] + "." + quotes
         # Handle multi-line docstrings
         else:
             lines = tokeninfo.string.splitlines()
@@ -124,14 +134,20 @@ def _treat_string(self, tokeninfo: tokenize.TokenInfo, indent_length: int) -> st
         return tokeninfo.string
 
 
-class StripWhitespacesFormatter(StringFormatter):
+class StripWhitespacesFormatter(StringAndQuotesFormatter):
     """Strip 1) docstring start, 2) docstring end and 3) end of line."""
 
     name = "strip-whitespaces"
 
-    def _treat_string(self, tokeninfo: tokenize.TokenInfo, indent_length: int) -> str:
+    def _treat_string(
+        self,
+        tokeninfo: tokenize.TokenInfo,
+        indent_length: int,
+        quotes: str,
+        quotes_length: Literal[1, 3],
+    ) -> str:
         """Strip whitespaces."""
-        quotes, lines = tokeninfo.string[:3], tokeninfo.string[3:-3].split("\n")
+        lines = tokeninfo.string[quotes_length:-quotes_length].split("\n")
         for index, line in enumerate(lines):
             if index == 0:  # pylint: disable=compare-to-zero
                 lines[index] = line.lstrip().rstrip()

tests/data/format/capitalization_first_letter/single_quote_docstring.py

@@ -0,0 +1,2 @@
+def func():
+    "a docstring."

tests/data/format/capitalization_first_letter/single_quote_docstring.py.out

@@ -0,0 +1,2 @@
+def func():
+    "A docstring."

tests/data/format/final_period/single_quote_docstring.py

@@ -0,0 +1,2 @@
+def func():
+    "A docstring"

tests/data/format/final_period/single_quote_docstring.py.out

@@ -0,0 +1,2 @@
+def func():
+    "A docstring."

tests/data/format/whitespace_stripper/single_quote_docstring.py

@@ -0,0 +1,2 @@
+def func():
+    " A docstring. "

tests/data/format/whitespace_stripper/single_quote_docstring.py.out

@@ -0,0 +1,2 @@
+def func():
+    "A docstring."