Fix fstrings not properly displaying when stringifying

dosisod · Feb 24, 2024 · 900f6b4 · 900f6b4
1 parent 12b7f84
commit 900f6b4
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 0 deletions.
diff --git a/refurb/checks/common.py b/refurb/checks/common.py
@@ -292,6 +292,49 @@ def is_type_none_call(node: Expression) -> bool:
     return False
 
 
+def get_fstring_parts(expr: Expression) -> list[tuple[bool, Expression, str]]:
+    match expr:
+        case CallExpr(
+            callee=MemberExpr(
+                expr=StrExpr(value="{:{}}"),
+                name="format",
+            ),
+            args=[arg, StrExpr(value=format_arg)],
+            arg_kinds=[ArgKind.ARG_POS, ArgKind.ARG_POS],
+        ):
+            return [(True, arg, format_arg)]
+
+        case CallExpr(
+            callee=MemberExpr(
+                expr=StrExpr(value=""),
+                name="join",
+            ),
+            args=[ListExpr(items=items)],
+            arg_kinds=[ArgKind.ARG_POS],
+        ):
+            exprs: list[tuple[bool, Expression, str]] = []
+
+            had_at_least_one_fstring_part = False
+
+            for item in items:
+                if isinstance(item, StrExpr):
+                    exprs.append((False, item, ""))
+
+                elif tmp := get_fstring_parts(item):
+                    had_at_least_one_fstring_part = True
+                    exprs.extend(tmp)
+
+                else:
+                    return []
+
+            if not had_at_least_one_fstring_part:
+                return []
+
+            return exprs
+
+    return []
+
+
 def stringify(node: Node) -> str:
     try:
         return _stringify(node)
@@ -352,6 +395,24 @@ def _stringify(node: Node) -> str:
             return f"({inner})"
 
         case CallExpr(arg_names=arg_names, arg_kinds=arg_kinds, args=args):
+            if fstring_parts := get_fstring_parts(node):
+                output = 'f"'
+
+                for is_format_arg, arg, fmt in fstring_parts:
+                    if not is_format_arg:
+                        assert isinstance(arg, StrExpr)
+
+                        output += arg.value
+
+                    elif fmt:
+                        output += f"{{{_stringify(arg)}:{fmt}}}"
+
+                    else:
+                        output += f"{{{_stringify(arg)}}}"
+
+                output += '"'
+                return output
+
             call_args: list[str] = []
 
             for arg_name, kind, arg in zip(arg_names, arg_kinds, args):

diff --git a/test/data/stringify.py b/test/data/stringify.py
@@ -18,3 +18,14 @@
 lambda x: x[:2:]
 lambda x: x[::]
 lambda x: x[:]
+
+# test fstring formatting
+_ = str(f"{123}")  # noqa: FURB183
+_ = str(f"{123:x}")
+_ = str(f"x{123}y")
+_ = str(f"x{123}y{456}z")
+_ = str(f"{'abc'}")  # noqa: FURB183
+
+# wont trigger string formatting
+_ = str("".join([""]))
+_ = str("".join(["", 1]))  # type: ignore
diff --git a/test/data/stringify.txt b/test/data/stringify.txt
@@ -13,3 +13,10 @@ test/data/stringify.py:17:1 [FURB118]: Replace `lambda x: x[::3]` with `operator
 test/data/stringify.py:18:1 [FURB118]: Replace `lambda x: x[:2]` with `operator.itemgetter(slice(None, 2))`
 test/data/stringify.py:19:1 [FURB118]: Replace `lambda x: x[:]` with `operator.itemgetter(slice(None, None))`
 test/data/stringify.py:20:1 [FURB118]: Replace `lambda x: x[:]` with `operator.itemgetter(slice(None, None))`
+test/data/stringify.py:23:5 [FURB123]: Replace `str(f"{123}")` with `f"{123}"`
+test/data/stringify.py:24:5 [FURB123]: Replace `str(f"{123:x}")` with `f"{123:x}"`
+test/data/stringify.py:25:5 [FURB123]: Replace `str(f"x{123}y")` with `f"x{123}y"`
+test/data/stringify.py:26:5 [FURB123]: Replace `str(f"x{123}y{456}z")` with `f"x{123}y{456}z"`
+test/data/stringify.py:27:5 [FURB123]: Replace `str(f"{"abc"}")` with `f"{"abc"}"`
+test/data/stringify.py:30:5 [FURB123]: Replace `str("".join([""]))` with `"".join([""])`
+test/data/stringify.py:31:5 [FURB123]: Replace `str("".join(["", 1]))` with `"".join(["", 1])`