Repository: asottile/pyupgrade

Issue:
unicode literal fixer `return'foo'`
Mistakenly does the following:

```diff
-    return'foo'
+    retrn'foo'
```

Discussion:

Patch:
diff --git a/pyupgrade.py b/pyupgrade.py
index a9c6cd5..0dcc2e0 100644
--- a/pyupgrade.py
+++ b/pyupgrade.py
@@ -490,6 +490,10 @@ def _fix_octal(s):
         return '0o' + s[1:]
 
 
+def _is_string_prefix(token):
+    return token.name == 'NAME' and set(token.src.lower()) <= set('bfru')
+
+
 def _fix_tokens(contents_text, py3_plus):
     remove_u_prefix = py3_plus or _imports_unicode_literals(contents_text)
 
@@ -503,7 +507,7 @@ def _fix_tokens(contents_text, py3_plus):
         elif token.name == 'STRING':
             # when a string prefix is not recognized, the tokenizer produces a
             # NAME token followed by a STRING token
-            if i > 0 and tokens[i - 1].name == 'NAME':
+            if i > 0 and _is_string_prefix(tokens[i - 1]):
                 tokens[i] = token._replace(src=tokens[i - 1].src + token.src)
                 tokens[i - 1] = tokens[i - 1]._replace(src='')
 


Test patch:
diff --git a/tests/pyupgrade_test.py b/tests/pyupgrade_test.py
index 9c6db3d..4c56411 100644
--- a/tests/pyupgrade_test.py
+++ b/tests/pyupgrade_test.py
@@ -323,6 +323,14 @@ def test_imports_unicode_literals(s, expected):
         ),
         # Regression: string containing newline
         ('"""with newline\n"""', True, '"""with newline\n"""'),
+        pytest.param(
+            'def f():\n'
+            '    return"foo"\n',
+            True,
+            'def f():\n'
+            '    return"foo"\n',
+            id='Regression: no space between return and string',
+        ),
     ),
 )
 def test_unicode_literals(s, py3_plus, expected):