Merge pull request #384 from asottile/reorg

reorganize tests and combine six_calls plugins

pyupgrade/_plugins/six_calls.py

@@ -2,18 +2,24 @@
 import functools
 import sys
 from typing import Iterable
+from typing import List
 from typing import Tuple
 from typing import Type
 
 from tokenize_rt import Offset
+from tokenize_rt import Token
 
 from pyupgrade._ast_helpers import ast_to_offset
 from pyupgrade._ast_helpers import has_starargs
 from pyupgrade._ast_helpers import is_name_attr
 from pyupgrade._data import register
 from pyupgrade._data import State
 from pyupgrade._data import TokenFunc
+from pyupgrade._string_helpers import is_ascii
 from pyupgrade._token_helpers import find_and_replace_call
+from pyupgrade._token_helpers import find_open_paren
+from pyupgrade._token_helpers import parse_call_args
+from pyupgrade._token_helpers import replace_call
 
 _EXPR_NEEDS_PARENS: Tuple[Type[ast.expr], ...] = (
     ast.Await, ast.BinOp, ast.BoolOp, ast.Compare, ast.GeneratorExp, ast.IfExp,
@@ -51,6 +57,17 @@
 RERAISE_3_TMPL = 'raise {args[1]}.with_traceback({args[2]})'
 
 
+def _fix_six_b(i: int, tokens: List[Token]) -> None:
+    j = find_open_paren(tokens, i)
+    if (
+            tokens[j + 1].name == 'STRING' and
+            is_ascii(tokens[j + 1].src) and
+            tokens[j + 2].src == ')'
+    ):
+        func_args, end = parse_call_args(tokens, j)
+        replace_call(tokens, i, end, func_args, 'b{args[0]}')
+
+
 @register(ast.Call)
 def visit_Call(
         state: State,
@@ -116,6 +133,20 @@ def visit_Call(
             template=SIX_INT2BYTE_TMPL,
         )
         yield ast_to_offset(node), func
+    elif (
+            state.settings.min_version >= (3,) and
+            is_name_attr(
+                node.func,
+                state.from_imports,
+                'six',
+                ('b', 'ensure_binary'),
+            ) and
+            not node.keywords and
+            not has_starargs(node) and
+            len(node.args) == 1 and
+            isinstance(node.args[0], ast.Str)
+    ):
+        yield ast_to_offset(node), _fix_six_b
     elif (
             isinstance(parent, ast.Expr) and
             is_name_attr(