Fix "string mode".

Previously parenthesis were (incorrectly) highlighted even in strings.

Now they are correctly ignored inside.

2 files changed, 56 insertions, 33 deletions

diff --git a/rainbow_parenthesis/__init__.py b/rainbow_parenthesis/__init__.py
index 1f5e2f7..e665bdb 100644
--- a/rainbow_parenthesis/__init__.py
+++ b/rainbow_parenthesis/__init__.py
@@ -80,39 +80,46 @@ def colorize(strm: io.TextIOBase) -> Generator[str | Colored, None, None]:
     depth = Stackpointer()
 
     while c := strm.read(1):
-        match c:
-            case '(':
-                yield color(depth(1), c)
-                paren += 1
-            case ')':
-                yield color(depth(-1), c)
-                paren -= 1
-            case '[':
-                yield color(depth(1), c)
-                brack += 1
-            case ']':
-                yield color(depth(-1), c)
-                brack -= 1
-            case '{':
-                yield color(depth(1), c)
-                brace += 1
-            case '}':
-                yield color(depth(-1), c)
-                brace -= 1
-            case "'":
-                if in_string == "'":
-                    in_string = False
-                else:
-                    in_string = "'"
+        match [in_string, c]:
+            case ['"', '"']:
+                # The string is ending
+                in_string = False
                 yield c
-            case '"':
-                if in_string == '"':
-                    in_string = False
-                else:
-                    in_string = '"'
+            case ["'", "'"]:
+                in_string = False
+                # The string is ending
                 yield c
-            case '\\':
-                yield c
-                yield strm.read(1)
-            case c:
+            case [False, _]:
+                # We are not in a string, emit the next character
+                # colored if it's a parenthesis (or similar), and
+                # plainly otherwise.
+                # Can also ether strings
+                match c:
+                    case '(':
+                        yield color(depth(1), c)
+                        paren += 1
+                    case ')':
+                        yield color(depth(-1), c)
+                        paren -= 1
+                    case '[':
+                        yield color(depth(1), c)
+                        brack += 1
+                    case ']':
+                        yield color(depth(-1), c)
+                        brack -= 1
+                    case '{':
+                        yield color(depth(1), c)
+                        brace += 1
+                    case '}':
+                        yield color(depth(-1), c)
+                        brace -= 1
+                    case "'" | '"':
+                        in_string = c
+                        yield c
+                    case '\\':
+                        yield c
+                        yield strm.read(1)
+                    case c:
+                        yield c
+            case [_, _]:
                 yield c
diff --git a/tests/test_string.py b/tests/test_string.py
index 5d3040b..1101d26 100644
--- a/tests/test_string.py
+++ b/tests/test_string.py
@@ -4,3 +4,19 @@ import io
 
 def test_escaped_parenthesis():
     assert list(r"\(") == list(colorize(io.StringIO(r"\(")))
+
+
+def test_escaped_parenthesis_in_string():
+    assert list(r"'\('") == list(colorize(io.StringIO(r"'\('")))
+
+
+def test_escaped_quot_in_string():
+    assert list(r"'\''") == list(colorize(io.StringIO(r"'\''")))
+
+
+def test_embedded_string():
+    assert list(colorize(io.StringIO("( '(' )"))) == [
+        Colored(depth=0, item='('),
+        *" '(' ",
+        Colored(depth=0, item=')'),
+    ]