From ca736a0107d1fe4710a2229e0e26c490350b6263 Mon Sep 17 00:00:00 2001
From: Anthony Sottile <asottile@umich.edu>
Date: Tue, 10 Nov 2020 15:56:55 -0800
Subject: [PATCH] fix rewrite causing syntax error when the first arg has
 newlines

---
 pyupgrade/_token_helpers.py | 18 ++++++++++++++++++
 tests/features/six_test.py  | 22 ++++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/pyupgrade/_token_helpers.py b/pyupgrade/_token_helpers.py
index 9669ffa8..44d9b575 100644
--- a/pyupgrade/_token_helpers.py
+++ b/pyupgrade/_token_helpers.py
@@ -396,6 +396,16 @@ def arg_str(tokens: List[Token], start: int, end: int) -> str:
     return tokens_to_src(tokens[start:end]).strip()
 
 
+def _arg_contains_newline(tokens: List[Token], start: int, end: int) -> bool:
+    while tokens[start].name in {'NL', 'NEWLINE', UNIMPORTANT_WS}:
+        start += 1
+    for i in range(start, end):
+        if tokens[i].name in {'NL', 'NEWLINE'}:
+            return True
+    else:
+        return False
+
+
 def replace_call(
         tokens: List[Token],
         start: int,
@@ -409,6 +419,14 @@ def replace_call(
     for paren in parens:
         arg_strs[paren] = f'({arg_strs[paren]})'
 
+    # there are a few edge cases which cause syntax errors when the first
+    # argument contains newlines (especially when moved outside of a natural
+    # contiunuation context)
+    if _arg_contains_newline(tokens, *args[0]) and 0 not in parens:
+        # this attempts to preserve more of the whitespace by using the
+        # original non-stripped argument string
+        arg_strs[0] = f'({tokens_to_src(tokens[slice(*args[0])])})'
+
     start_rest = args[0][1] + 1
     while (
             start_rest < end and
diff --git a/tests/features/six_test.py b/tests/features/six_test.py
index 71d41c8b..2b87b8ff 100644
--- a/tests/features/six_test.py
+++ b/tests/features/six_test.py
@@ -347,6 +347,28 @@ def test_fix_six_noop(s):
             '(x < y).values()',
             id='needs parentehsizing for Compare',
         ),
+        pytest.param(
+            'x = six.itervalues(\n'
+            '    # comment\n'
+            '    x\n'
+            ')',
+            'x = (\n'
+            '    # comment\n'
+            '    x\n'
+            ').values()',
+            id='multiline first argument with comment',
+        ),
+        pytest.param(
+            'x = six.itervalues(\n'
+            '    # comment\n'
+            '    x,\n'
+            ')',
+            # TODO: ideally this would preserve whitespace better
+            'x = (\n'
+            '    # comment\n'
+            '    x).values()',
+            id='multiline first argument with comment, trailing comma',
+        ),
     ),
 )
 def test_fix_six(s, expected):