handle redundant t mode

asottile · asottile · commit 981bbc6269d2 · 2022-10-06T16:02:38.000-07:00
diff --git a/README.md b/README.md
@@ -487,6 +487,8 @@ Note that `if` blocks without an `else` will not be rewritten as it could introd
 +open("foo")
 -open("f", "r", encoding="UTF-8")
 +open("f", encoding="UTF-8")
+-open("f", "wt")
++open("f", "w")
 ```
 
 
diff --git a/pyupgrade/_plugins/open_mode.py b/pyupgrade/_plugins/open_mode.py
@@ -2,6 +2,7 @@
 
 import ast
 import functools
+import itertools
 from typing import Iterable
 from typing import NamedTuple
 
@@ -18,10 +19,20 @@
 from pyupgrade._token_helpers import find_open_paren
 from pyupgrade._token_helpers import parse_call_args
 
-U_MODE_REMOVE = frozenset(('U', 'Ur', 'rU', 'r', 'rt', 'tr'))
-U_MODE_REPLACE_R = frozenset(('Ub', 'bU'))
-U_MODE_REMOVE_U = frozenset(('rUb', 'Urb', 'rbU', 'Ubr', 'bUr', 'brU'))
-U_MODE_REPLACE = U_MODE_REPLACE_R | U_MODE_REMOVE_U
+
+def _plus(args: tuple[str, ...]) -> tuple[str, ...]:
+    return args + tuple(f'{arg}+' for arg in args)
+
+
+def _permute(*args: str) -> tuple[str, ...]:
+    return tuple(''.join(p) for s in args for p in itertools.permutations(s))
+
+
+MODE_REMOVE = frozenset(_permute('U', 'r', 'rU', 'rt'))
+MODE_REPLACE_R = frozenset(_permute('Ub'))
+MODE_REMOVE_T = frozenset(_plus(_permute('at', 'rt', 'wt', 'xt')))
+MODE_REMOVE_U = frozenset(_permute('rUb'))
+MODE_REPLACE = MODE_REPLACE_R | MODE_REMOVE_T | MODE_REMOVE_U
 
 
 class FunctionArg(NamedTuple):
@@ -35,12 +46,15 @@ def _fix_open_mode(i: int, tokens: list[Token], *, arg_idx: int) -> None:
     mode = tokens_to_src(tokens[slice(*func_args[arg_idx])])
     mode_stripped = mode.split('=')[-1]
     mode_stripped = ast.literal_eval(mode_stripped.strip())
-    if mode_stripped in U_MODE_REMOVE:
+    if mode_stripped in MODE_REMOVE:
         delete_argument(arg_idx, tokens, func_args)
-    elif mode_stripped in U_MODE_REPLACE_R:
+    elif mode_stripped in MODE_REPLACE_R:
         new_mode = mode.replace('U', 'r')
         tokens[slice(*func_args[arg_idx])] = [Token('SRC', new_mode)]
-    elif mode_stripped in U_MODE_REMOVE_U:
+    elif mode_stripped in MODE_REMOVE_T:
+        new_mode = mode.replace('t', '')
+        tokens[slice(*func_args[arg_idx])] = [Token('SRC', new_mode)]
+    elif mode_stripped in MODE_REMOVE_U:
         new_mode = mode.replace('U', '')
         tokens[slice(*func_args[arg_idx])] = [Token('SRC', new_mode)]
     else:
@@ -69,13 +83,10 @@ def visit_Call(
     ):
         if len(node.args) >= 2 and isinstance(node.args[1], ast.Str):
             if (
-                node.args[1].s in U_MODE_REPLACE or
-                (len(node.args) == 2 and node.args[1].s in U_MODE_REMOVE)
+                node.args[1].s in MODE_REPLACE or
+                (len(node.args) == 2 and node.args[1].s in MODE_REMOVE)
             ):
-                func = functools.partial(
-                    _fix_open_mode,
-                    arg_idx=1,
-                )
+                func = functools.partial(_fix_open_mode, arg_idx=1)
                 yield ast_to_offset(node), func
         elif node.keywords and (len(node.keywords) + len(node.args) > 1):
             mode = next(
@@ -90,8 +101,8 @@ def visit_Call(
                 mode is not None and
                 isinstance(mode.value, ast.Str) and
                 (
-                    mode.value.s in U_MODE_REMOVE or
-                    mode.value.s in U_MODE_REPLACE
+                    mode.value.s in MODE_REMOVE or
+                    mode.value.s in MODE_REPLACE
                 )
             ):
                 func = functools.partial(
diff --git a/tests/features/open_mode_test.py b/tests/features/open_mode_test.py
@@ -4,6 +4,18 @@
 
 from pyupgrade._data import Settings
 from pyupgrade._main import _fix_plugins
+from pyupgrade._plugins.open_mode import _permute
+from pyupgrade._plugins.open_mode import _plus
+
+
+def test_plus():
+    assert _plus(('a',)) == ('a', 'a+')
+    assert _plus(('a', 'b')) == ('a', 'b', 'a+', 'b+')
+
+
+def test_permute():
+    assert _permute('ab') == ('ab', 'ba')
+    assert _permute('abc') == ('abc', 'acb', 'bac', 'bca', 'cab', 'cba')
 
 
 @pytest.mark.parametrize(
@@ -18,8 +30,6 @@
         'open("foo", qux="r")',
         'open("foo", 3)',
         'open(mode="r")',
-        # TODO: could maybe be rewritten to remove t?
-        'open("foo", "wt")',
         # don't remove this, they meant to use `encoding=`
         'open("foo", "r", "utf-8")',
     ),
@@ -70,6 +80,11 @@ def test_fix_open_mode_noop(s):
             'open("foo")',
             id='io.open also rewrites modes in a single pass',
         ),
+        ('open("foo", "wt")', 'open("foo", "w")'),
+        ('open("foo", "xt")', 'open("foo", "x")'),
+        ('open("foo", "at")', 'open("foo", "a")'),
+        ('open("foo", "wt+")', 'open("foo", "w+")'),
+        ('open("foo", "rt+")', 'open("foo", "r+")'),
     ),
 )
 def test_fix_open_mode(s, expected):
