Fix incorrect error message for FURB122

dosisod · Mar 1, 2024 · 6b9035f · 6b9035f
1 parent 60bb2f6
commit 6b9035f
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 20 deletions.
diff --git a/docs/checks.md b/docs/checks.md
@@ -583,13 +583,6 @@ with open("file") as f:
     f.writelines(lines)
 ```
 
-Note: If you have a more complex expression then just `lines`, you may
-need to use a list comprehension instead. For example:
-
-```python
-f.writelines(f"{line}\n" for line in lines)
-```
-
 ## FURB123: `no-redundant-cast`
 
 Categories: `readability`

diff --git a/refurb/checks/builtin/writelines.py b/refurb/checks/builtin/writelines.py
@@ -39,13 +39,6 @@ class ErrorInfo(Error):
     with open("file") as f:
         f.writelines(lines)
     ```
-
-    Note: If you have a more complex expression then just `lines`, you may
-    need to use a list comprehension instead. For example:
-
-    ```
-    f.writelines(f"{line}\n" for line in lines)
-    ```
     """
 
     name = "use-writelines"
@@ -65,7 +58,7 @@ def check(node: WithStmt, errors: list[Error]) -> None:
             body=Block(
                 body=[
                     ForStmt(
-                        index=NameExpr(),
+                        index=NameExpr() as for_target,
                         expr=source,
                         body=Block(
                             body=[
@@ -74,7 +67,8 @@ def check(node: WithStmt, errors: list[Error]) -> None:
                                         callee=MemberExpr(
                                             expr=NameExpr() as write_base,
                                             name="write",
-                                        )
+                                        ),
+                                        args=[write_arg],
                                     )
                                 )
                             ]
@@ -86,8 +80,12 @@ def check(node: WithStmt, errors: list[Error]) -> None:
             ),
         ) if is_file_object(f) and is_equivalent(f, write_base):
             old = stringify(for_stmt)
-            new = f"{stringify(f)}.writelines({stringify(source)})"
 
-            msg = f"Replace `{old}` with `{new}`"
+            if is_equivalent(for_target, write_arg):
+                new = stringify(source)
+            else:
+                new = f"{stringify(write_arg)} for {stringify(for_target)} in {stringify(source)}"
+
+            msg = f"Replace `{old}` with `{stringify(f)}.writelines({new})`"
 
             errors.append(ErrorInfo.from_node(for_stmt, msg))
diff --git a/test/data/err_122.py b/test/data/err_122.py
@@ -44,3 +44,8 @@ async def func():
     with open("file") as f:
         async for line in lines:  # type: ignore
             f.write(line)
+
+
+with open("file") as f:
+    for line in lines:
+        f.write()  # type: ignore
diff --git a/test/data/err_122.txt b/test/data/err_122.txt
@@ -1,3 +1,3 @@
 test/data/err_122.py:6:5 [FURB122]: Replace `for line in lines: f.write(line)` with `f.writelines(lines)`
-test/data/err_122.py:11:5 [FURB122]: Replace `for line in lines: f.write(line.encode())` with `f.writelines(lines)`
-test/data/err_122.py:16:5 [FURB122]: Replace `for line in lines: f.write(line.upper())` with `f.writelines(lines)`
+test/data/err_122.py:11:5 [FURB122]: Replace `for line in lines: f.write(line.encode())` with `f.writelines(line.encode() for line in lines)`
+test/data/err_122.py:16:5 [FURB122]: Replace `for line in lines: f.write(line.upper())` with `f.writelines(line.upper() for line in lines)`