py: fix bugs related to switching a snapshot between string and repr (#385)

Author: nedtwigg

Diff for: python/selfie-lib/selfie_lib/Literals.py

@@ -44,36 +44,22 @@ def parse(self, string: str, language: Language) -> T:
 PADDING_SIZE = len(str(MAX_RAW_NUMBER)) - 1
 
 
-class LiteralInt(LiteralFormat[int]):
-    def _encode_underscores(
-        self, buffer: io.StringIO, value: int, language: Language
-    ) -> io.StringIO:
-        if value >= MAX_RAW_NUMBER:
-            mod = value % MAX_RAW_NUMBER
-            left_padding = PADDING_SIZE - len(str(mod))
-            self._encode_underscores(buffer, value // MAX_RAW_NUMBER, language)
-            buffer.write("_")
-            buffer.write("0" * left_padding)
-            buffer.write(str(mod))
-            return buffer
-        elif value < 0:
-            buffer.write("-")
-            self._encode_underscores(buffer, abs(value), language)
-            return buffer
-        else:
-            buffer.write(str(value))
-            return buffer
-
-    def encode(
-        self,
-        value: int,
-        language: Language,
-        encoding_policy: EscapeLeadingWhitespace,  # noqa: ARG002
-    ) -> str:
-        return self._encode_underscores(io.StringIO(), value, language).getvalue()
-
-    def parse(self, string: str, language: Language) -> int:  # noqa: ARG002
-        return int(string.replace("_", ""))
+def _encode_int_underscores(buffer: io.StringIO, value: int) -> str:
+    if value >= MAX_RAW_NUMBER:
+        mod = value % MAX_RAW_NUMBER
+        left_padding = PADDING_SIZE - len(str(mod))
+        _encode_int_underscores(buffer, value // MAX_RAW_NUMBER)
+        buffer.write("_")
+        buffer.write("0" * left_padding)
+        buffer.write(str(mod))
+        return buffer.getvalue()
+    elif value < 0:
+        buffer.write("-")
+        _encode_int_underscores(buffer, abs(value))
+        return buffer.getvalue()
+    else:
+        buffer.write(str(value))
+        return buffer.getvalue()
 
 
 TRIPLE_QUOTE = '"""'
@@ -252,15 +238,18 @@ def handle_escape_sequences(line: str) -> str:
 
 class LiteralRepr(LiteralFormat[Any]):
     def encode(
-        self, value: Any, language: Language, encoding_policy: EscapeLeadingWhitespace
+        self,
+        value: Any,
+        language: Language,  # noqa: ARG002
+        encoding_policy: EscapeLeadingWhitespace,  # noqa: ARG002
     ) -> str:
         if isinstance(value, int):
-            return LiteralInt().encode(value, language, encoding_policy)
+            return _encode_int_underscores(io.StringIO(), value)
         else:
             return repr(value)
 
-    def parse(self, string: str, language: Language) -> Any:  # noqa: ARG002
-        return eval(string)
+    def parse(self, string: str, language: Language) -> Any:
+        raise NotImplementedError
 
 
 class TodoStub(Enum):

Diff for: python/selfie-lib/selfie_lib/SelfieImplementations.py

@@ -158,11 +158,6 @@ def to_be_TODO(self, _: Any = None) -> str:
 
     def to_be(self, expected: str) -> str:
         actual_string = self.__actual()
-
-        # Check if expected is a string
-        if not isinstance(expected, str):
-            raise TypeError("Expected value must be a string")
-
         if actual_string == expected:
             return _checkSrc(actual_string)
         else:

Diff for: python/selfie-lib/selfie_lib/SourceFile.py

@@ -1,7 +1,7 @@
 from typing import Any
 
 from .EscapeLeadingWhitespace import EscapeLeadingWhitespace
-from .Literals import Language, LiteralFormat, LiteralValue
+from .Literals import Language, LiteralFormat, LiteralRepr, LiteralValue
 from .Slice import Slice
 
 
@@ -78,21 +78,25 @@ def _get_arg(self):
 
         def set_literal_and_get_newline_delta(self, literal_value: LiteralValue) -> int:
             encoded = literal_value.format.encode(
-                literal_value.actual, self.__language, self.__escape_leading_whitespace
+                literal_value.actual,
+                self.__language,
+                self.__escape_leading_whitespace,
             )
-            round_tripped = literal_value.format.parse(encoded, self.__language)
-            if round_tripped != literal_value.actual:
-                raise ValueError(
-                    f"There is an error in {literal_value.format.__class__.__name__}, "
-                    "the following value isn't round tripping.\n"
-                    f"Please report this error and the data below at "
-                    "https://github.com/diffplug/selfie/issues/new\n"
-                    f"```\n"
-                    f"ORIGINAL\n{literal_value.actual}\n"
-                    f"ROUNDTRIPPED\n{round_tripped}\n"
-                    f"ENCODED ORIGINAL\n{encoded}\n"
-                    f"```\n"
-                )
+            if not isinstance(literal_value.format, LiteralRepr):
+                # we don't roundtrip LiteralRepr because `eval` is dangerous
+                round_tripped = literal_value.format.parse(encoded, self.__language)
+                if round_tripped != literal_value.actual:
+                    raise ValueError(
+                        f"There is an error in {literal_value.format.__class__.__name__}, "
+                        "the following value isn't round tripping.\n"
+                        f"Please report this error and the data below at "
+                        "https://github.com/diffplug/selfie/issues/new\n"
+                        f"```\n"
+                        f"ORIGINAL\n{literal_value.actual}\n"
+                        f"ROUNDTRIPPED\n{round_tripped}\n"
+                        f"ENCODED ORIGINAL\n{encoded}\n"
+                        f"```\n"
+                    )
             existing_newlines = self.__function_call_plus_arg.count("\n")
             new_newlines = encoded.count("\n")
             self.__parent._content_slice = Slice(  # noqa: SLF001
@@ -153,48 +157,98 @@ def parse_to_be_like(self, line_one_indexed: int) -> ToBeLiteral:
                     f"on line {line_one_indexed}"
                 )
 
-        end_arg = -1
-        end_paren = 0
         if self._content_slice[arg_start] == '"':
-            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)
+            (end_paren, end_arg) = self._parse_string(
+                line_one_indexed, arg_start, dot_fun_open_paren
+            )
+        else:
+            (end_paren, end_arg) = self._parse_code(
+                line_one_indexed, arg_start, dot_fun_open_paren
+            )
+        return self.ToBeLiteral(
+            self,
+            dot_fun_open_paren.replace("_TODO", ""),
+            self._content_slice.subSequence(dot_function_call, end_paren + 1),
+            self._content_slice.subSequence(arg_start, end_arg),
+            self.__language,
+            self.__escape_leading_whitespace,
+        )
+
+    def _parse_code(
+        self,
+        line_one_indexed: int,
+        arg_start: int,
+        dot_fun_open_paren: str,
+    ):
+        # Initialize variables
+        parenthesis_count = 1
+        string_delimiter = None
+
+        # Iterate through the characters starting from the given index
+        for i in range(arg_start, len(self._content_slice)):
+            char = self._content_slice[i]
+
+            # Check if we are entering or leaving a string
+            if char in ["'", '"'] and self._content_slice[i - 1] != "\\":
+                if not string_delimiter:
+                    string_delimiter = char
+                elif char == string_delimiter:
+                    string_delimiter = None
+
+            # Skip characters inside strings
+            if string_delimiter:
+                continue
+
+            # Count parentheses
+            if char == "(":
+                parenthesis_count += 1
+            elif char == ")":
+                parenthesis_count -= 1
+
+            # If all parentheses are closed, return the current index
+            if parenthesis_count == 0:
+                end_paren = i
+                end_arg = i - 1
+                return (end_paren, end_arg)
+        # else ...
+        raise AssertionError(
+            f"Appears to be an unclosed function call `{dot_fun_open_paren}` "
+            f"starting at line {line_one_indexed}"
+        )
+
+    def _parse_string(
+        self,
+        line_one_indexed: int,
+        arg_start: int,
+        dot_fun_open_paren: str,
+    ):
+        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)
+            )
+            if end_arg == -1:
+                raise AssertionError(
+                    f"Appears to be an unclosed multiline string literal `{self.TRIPLE_QUOTE}` "
+                    f"on line {line_one_indexed}"
                 )
-                if end_arg == -1:
-                    raise AssertionError(
-                        f"Appears to be an unclosed multiline string literal `{self.TRIPLE_QUOTE}` "
-                        f"on line {line_one_indexed}"
-                    )
-                else:
-                    end_arg += len(self.TRIPLE_QUOTE)
-                    end_paren = end_arg
             else:
-                end_arg = arg_start + 1
-                while (
-                    self._content_slice[end_arg] != '"'
-                    or self._content_slice[end_arg - 1] == "\\"
-                ):
-                    end_arg += 1
-                    if end_arg == self._content_slice.__len__():
-                        raise AssertionError(
-                            f'Appears to be an unclosed string literal `"` '
-                            f"on line {line_one_indexed}"
-                        )
-                end_arg += 1
+                end_arg += len(self.TRIPLE_QUOTE)
                 end_paren = end_arg
         else:
-            end_arg = arg_start
-            while not self._content_slice[end_arg].isspace():
-                if self._content_slice[end_arg] == ")":
-                    break
+            end_arg = arg_start + 1
+            while (
+                self._content_slice[end_arg] != '"'
+                or self._content_slice[end_arg - 1] == "\\"
+            ):
                 end_arg += 1
                 if end_arg == self._content_slice.__len__():
                     raise AssertionError(
-                        f"Appears to be an unclosed numeric literal "
+                        f'Appears to be an unclosed string literal `"` '
                         f"on line {line_one_indexed}"
                     )
+            end_arg += 1
             end_paren = end_arg
         while self._content_slice[end_paren] != ")":
             if not self._content_slice[end_paren].isspace():
@@ -210,14 +264,7 @@ def parse_to_be_like(self, line_one_indexed: int) -> ToBeLiteral:
                     f"Appears to be an unclosed function call `{dot_fun_open_paren}` "
                     f"starting at line {line_one_indexed}"
                 )
-        return self.ToBeLiteral(
-            self,
-            dot_fun_open_paren.replace("_TODO", ""),
-            self._content_slice.subSequence(dot_function_call, end_paren + 1),
-            self._content_slice.subSequence(arg_start, end_arg),
-            self.__language,
-            self.__escape_leading_whitespace,
-        )
+        return (end_paren, end_arg)
 
 
 TO_BE_LIKES = [

Diff for: python/selfie-lib/selfie_lib/WriteTracker.py

@@ -7,7 +7,7 @@
 from typing import Dict, Generic, List, Optional, TypeVar, cast
 
 from .FS import FS
-from .Literals import LiteralTodoStub, LiteralValue, TodoStub
+from .Literals import LiteralString, LiteralTodoStub, LiteralValue, TodoStub
 from .SourceFile import SourceFile
 from .TypedPath import TypedPath
 
@@ -175,7 +175,11 @@ def record(
 
         file = layout.sourcefile_for_call(call.location)
 
-        if snapshot.expected is not None:
+        if (
+            snapshot.expected is not None
+            and isinstance(snapshot.expected, str)
+            and isinstance(snapshot.format, LiteralString)
+        ):
             content = SourceFile(file.name, layout.fs.file_read(file))
             try:
                 snapshot = cast(LiteralValue, snapshot)

Diff for: python/selfie-lib/tests/LiteralInt_test.py

@@ -1,17 +1,11 @@
-from selfie_lib.EscapeLeadingWhitespace import EscapeLeadingWhitespace
-from selfie_lib.Literals import Language, LiteralInt
+import io
 
-
-def _encode(value: int, expected: str):
-    literal_int = LiteralInt()
-    actual = literal_int.encode(value, Language.PYTHON, EscapeLeadingWhitespace.NEVER)
-    assert actual == expected, f"Expected '{expected}', but got '{actual}'"
+from selfie_lib.Literals import _encode_int_underscores
 
 
-def _decode(value: str, expected: int):
-    literal_int = LiteralInt()
-    actual = literal_int.parse(value, Language.PYTHON)
-    assert actual == expected, f"Expected '{expected}', but got '{actual}'"
+def _encode(value: int, expected: str):
+    actual = _encode_int_underscores(io.StringIO(), value)
+    assert actual == expected
 
 
 class TestLiteralInt:
@@ -35,18 +29,3 @@ def test_encode(self):
         ]
         for value, expected in test_cases:
             _encode(value, expected)
-
-    def test_decode(self):
-        test_cases = [
-            ("0", 0),
-            ("1", 1),
-            ("-1", -1),
-            ("999", 999),
-            ("9_99", 999),
-            ("9_9_9", 999),
-            ("-999", -999),
-            ("-9_99", -999),
-            ("-9_9_9", -999),
-        ]
-        for value, expected in test_cases:
-            _decode(value, expected)