Fix f str = handling

coconut/compiler/compiler.py

@@ -105,6 +105,7 @@
     get_clock_time,
     get_name,
     assert_remove_prefix,
+    assert_remove_suffix,
     dictset,
     noop_ctx,
 )
@@ -4335,11 +4336,11 @@ def f_string_handle(self, original, loc, tokens):
 
         # handle Python 3.8 f string = specifier
         for i, expr in enumerate(exprs):
-            if expr.endswith("="):
+            expr_rstrip = expr.rstrip()
+            if expr_rstrip.endswith("="):
                 before = string_parts[i]
-                internal_assert(before[-1] == "{", "invalid format string split", (string_parts, exprs))
-                string_parts[i] = before[:-1] + expr + "{"
-                exprs[i] = expr[:-1]
+                string_parts[i] = assert_remove_prefix(before, "{") + expr + "{"
+                exprs[i] = assert_remove_suffix(expr_rstrip, "=")
 
         # compile Coconut expressions
         compiled_exprs = []

coconut/tests/src/cocotest/agnostic/primary_2.coco

@@ -459,6 +459,8 @@ def primary_test_2() -> bool:
     assert py_min(3, 4) == 3 == py_max(2, 3)
     assert len(zip()) == 0 == len(zip_longest())  # type: ignore
     assert CoconutWarning `issubclass` Warning
+    a = b = 2
+    assert f"{a + b = }" == "a + b = 4"
 
     with process_map.multiple_sequential_calls():  # type: ignore
         assert map((+), range(3), range(4)$[:-1], strict=True) |> list == [0, 2, 4] == process_map((+), range(3), range(4)$[:-1], strict=True) |> list  # type: ignore

coconut/util.py

@@ -291,6 +291,19 @@ def assert_remove_prefix(inputstr, prefix, allow_no_prefix=False):
 remove_prefix = partial(assert_remove_prefix, allow_no_prefix=True)
 
 
+def assert_remove_suffix(inputstr, suffix, allow_no_suffix=False):
+    """Remove prefix asserting that inputstr starts with it."""
+    assert suffix, suffix
+    if not allow_no_suffix:
+        assert inputstr.endswith(suffix), inputstr
+    elif not inputstr.endswith(suffix):
+        return inputstr
+    return inputstr[:-len(suffix)]
+
+
+remove_suffix = partial(assert_remove_suffix, allow_no_suffix=True)
+
+
 def ensure_dir(dirpath, logger=None):
     """Ensure that a directory exists."""
     if not os.path.exists(dirpath):