py: fixup multiline string literal roundtripping (#384)

diffplug · May 8, 2024 · 107bbee · 107bbee
2 parents 273b6c0 + e962ad1
commit 107bbee
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 5 deletions.
diff --git a/python/selfie-lib/selfie_lib/Literals.py b/python/selfie-lib/selfie_lib/Literals.py
@@ -160,7 +160,7 @@ def protect_trailing_whitespace(line: str) -> str:
             for line in lines
         )
 
-        return f"{TRIPLE_QUOTE}\n{protect_whitespace}{TRIPLE_QUOTE}"
+        return f"{TRIPLE_QUOTE}{protect_whitespace}{TRIPLE_QUOTE}"
 
     _char_literal_pattern = re.compile(r"""\{'(\\?.)'\}""")
 

diff --git a/python/selfie-lib/selfie_lib/Slice.py b/python/selfie-lib/selfie_lib/Slice.py
@@ -89,3 +89,8 @@ def count(self, char: str) -> int:
 
     def baseLineAtOffset(self, index: int) -> int:
         return 1 + Slice(self.base, 0, index).count("\n")
+
+    def starts_with(self, prefix: str) -> bool:
+        if len(prefix) > len(self):
+            return False
+        return all(self[i] == prefix[i] for i in range(len(prefix)))
diff --git a/python/selfie-lib/selfie_lib/SourceFile.py b/python/selfie-lib/selfie_lib/SourceFile.py
@@ -156,7 +156,9 @@ def parse_to_be_like(self, line_one_indexed: int) -> ToBeLiteral:
         end_arg = -1
         end_paren = 0
         if self._content_slice[arg_start] == '"':
-            if self._content_slice[arg_start].startswith(self.TRIPLE_QUOTE):
+            if self._content_slice.subSequence(
+                arg_start, len(self._content_slice)
+            ).starts_with(self.TRIPLE_QUOTE):
                 end_arg = self._content_slice.indexOf(
                     self.TRIPLE_QUOTE, arg_start + len(self.TRIPLE_QUOTE)
                 )

diff --git a/python/selfie-lib/tests/LiteralString_test.py b/python/selfie-lib/tests/LiteralString_test.py
@@ -26,11 +26,11 @@ def test_encode_single_with_dollars(self, value, expected):
     @pytest.mark.parametrize(
         ("value", "expected"),
         [
-            ("1", "'''\n1'''"),
-            ("\\", "'''\n\\\\'''"),
+            ("1", "'''1'''"),
+            ("\\", "'''\\\\'''"),
             (
                 "  leading\ntrailing  ",
-                "'''\n" + "  leading\n" + "trailing \\u0020'''",
+                "'''  leading\ntrailing \\u0020'''",
             ),
         ],
     )