Skip to content

Commit

Permalink
rewrite bytes(b'foo') to b'foo'
Browse files Browse the repository at this point in the history
  • Loading branch information
jkittner committed Jun 6, 2022
1 parent 54356ff commit 798700f
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 4 deletions.
19 changes: 16 additions & 3 deletions pyupgrade/_plugins/native_literals.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import ast
import functools
from typing import Iterable

from tokenize_rt import Offset
Expand All @@ -19,15 +20,15 @@
SIX_NATIVE_STR = frozenset(('ensure_str', 'ensure_text', 'text_type'))


def _fix_native_str(i: int, tokens: list[Token]) -> None:
def _fix_literal(i: int, tokens: list[Token], *, empty: str) -> None:
j = find_open_paren(tokens, i)
func_args, end = parse_call_args(tokens, j)
if any(tok.name == 'NL' for tok in tokens[i:end]):
return
if func_args:
replace_call(tokens, i, end, func_args, '{args[0]}')
else:
tokens[i:end] = [tokens[i]._replace(name='STRING', src="''")]
tokens[i:end] = [tokens[i]._replace(name='STRING', src=empty)]


def is_a_native_literal_call(
Expand Down Expand Up @@ -58,4 +59,16 @@ def visit_Call(
state.settings.min_version >= (3,) and
is_a_native_literal_call(node, state.from_imports)
):
yield ast_to_offset(node), _fix_native_str
func = functools.partial(_fix_literal, empty="''")
yield ast_to_offset(node), func
elif (
state.settings.min_version >= (3,) and
isinstance(node.func, ast.Name) and node.func.id == 'bytes' and
not node.keywords and not has_starargs(node) and
(
len(node.args) == 0 or
(len(node.args) == 1 and isinstance(node.args[0], ast.Bytes))
)
):
func = functools.partial(_fix_literal, empty="b''")
yield ast_to_offset(node), func
2 changes: 1 addition & 1 deletion pyupgrade/_token_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -420,7 +420,7 @@ def replace_call(

# 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)
# continuation 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
Expand Down
8 changes: 8 additions & 0 deletions tests/features/native_literals_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,11 @@
'str(*a)', 'str("foo", *a)',
'str(**k)', 'str("foo", **k)',
'str("foo", encoding="UTF-8")',
'bytes("foo", encoding="UTF-8")',
'bytes(b"foo"\nb"bar")',
'bytes("foo"\n"bar")',
'bytes(*a)', 'bytes("foo", *a)',
'bytes("foo", **a)',
),
)
def test_fix_native_literals_noop(s):
Expand All @@ -38,6 +43,9 @@ def test_fix_native_literals_noop(s):
id='from import of rewritten name',
),
('bytes()', "b''"),
('bytes(b"foo")', 'b"foo"'),
('bytes(b"""\nfoo""")', 'b"""\nfoo"""'),
),
)
def test_fix_native_literals(s, expected):
Expand Down

0 comments on commit 798700f

Please sign in to comment.