From 6b9035fa67f0651b21dabdb08d0cf03a8fca9010 Mon Sep 17 00:00:00 2001
From: dosisod <39638017+dosisod@users.noreply.github.com>
Date: Thu, 29 Feb 2024 23:09:31 -0800
Subject: [PATCH] Fix incorrect error message for FURB122

---
 docs/checks.md                      |  7 -------
 refurb/checks/builtin/writelines.py | 20 +++++++++-----------
 test/data/err_122.py                |  5 +++++
 test/data/err_122.txt               |  4 ++--
 4 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/docs/checks.md b/docs/checks.md
index 9a88c31..946bd16 100644
--- 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
index 6dd881b..21ee38b 100644
--- 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
index e7bc6b5..c6d4f2b 100644
--- 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
index 8b9838e..1546029 100644
--- 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)`