Repository: asottile/pyupgrade

Issue:
six: ensure_str / ensure_binary / ensure_text of a literal string
while we generally can't fix the ensure_* variants, we can for string literals (as we're already doing for `str('...')` for example)

Discussion:

Patch:
diff --git a/README.md b/README.md
index eeb1a5f..7a7a469 100644
--- a/README.md
+++ b/README.md
@@ -331,6 +331,11 @@ six.get_function_globals(fn)            # fn.__globals__
 six.assertCountEqual(self, a1, a2)      # self.assertCountEqual(a1, a2)
 six.assertRaisesRegex(self, e, r, fn)   # self.assertRaisesRegex(e, r, fn)
 six.assertRegex(self, s, r)             # self.assertRegex(s, r)
+
+# note: only for *literals*
+six.ensure_binary('...')                # b'...'
+six.ensure_str('...')                   # '...'
+six.ensure_text('...')                  # '...'
 ```
 
 ### `open` alias
diff --git a/pyupgrade.py b/pyupgrade.py
index 07ed723..d110895 100644
--- a/pyupgrade.py
+++ b/pyupgrade.py
@@ -1420,7 +1420,7 @@ class FindPy3Plus(ast.NodeVisitor):
             # _is_six() enforces this
             assert isinstance(arg, (ast.Name, ast.Attribute))
             self.six_type_ctx[_ast_to_offset(node.args[1])] = arg
-        elif self._is_six(node.func, ('b',)):
+        elif self._is_six(node.func, ('b', 'ensure_binary')):
             self.six_b.add(_ast_to_offset(node))
         elif self._is_six(node.func, SIX_CALLS) and not _starargs(node):
             self.six_calls[_ast_to_offset(node)] = node
@@ -1457,8 +1457,10 @@ class FindPy3Plus(ast.NodeVisitor):
         ):
             self.super_calls[_ast_to_offset(node)] = node
         elif (
-                isinstance(node.func, ast.Name) and
-                node.func.id == 'str' and
+                (
+                    self._is_six(node.func, ('ensure_str', 'ensure_text')) or
+                    isinstance(node.func, ast.Name) and node.func.id == 'str'
+                ) and
                 not node.keywords and
                 not _starargs(node) and
                 (


Test patch:
diff --git a/tests/native_literals_test.py b/tests/native_literals_test.py
index a82820a..7b6c8f0 100644
--- a/tests/native_literals_test.py
+++ b/tests/native_literals_test.py
@@ -27,6 +27,8 @@ def test_fix_native_literals_noop(s):
         ('str()', "''"),
         ('str("foo")', '"foo"'),
         ('str("""\nfoo""")', '"""\nfoo"""'),
+        ('six.ensure_str("foo")', '"foo"'),
+        ('six.ensure_text("foo")', '"foo"'),
     ),
 )
 def test_fix_native_literals(s, expected):
diff --git a/tests/six_test.py b/tests/six_test.py
index f840c3f..83ce558 100644
--- a/tests/six_test.py
+++ b/tests/six_test.py
@@ -98,6 +98,10 @@ def test_fix_six_noop(s):
             r'six.b("\x12\xef")',
             r'b"\x12\xef"',
         ),
+        (
+            'six.ensure_binary("foo")',
+            'b"foo"',
+        ),
         (
             'from six import b\n\n' r'b("\x12\xef")',
             'from six import b\n\n' r'b"\x12\xef"',