Skip to content

Commit

Permalink
rewrite six.moves.range
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed Jul 23, 2021
1 parent ca5d7df commit 943891c
Show file tree
Hide file tree
Showing 4 changed files with 47 additions and 0 deletions.
2 changes: 2 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -308,6 +308,8 @@ six.exec_(c, g, l) # exec(c, g, l)
six.advance_iterator(it) # next(it)
six.next(it) # next(it)
six.callable(x) # callable(x)
six.moves.range(x) # range(x)
six.moves.xrange(x) # range(x)

from six import text_type
text_type # str
Expand Down
1 change: 1 addition & 0 deletions pyupgrade/_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ class State(NamedTuple):
'mmap',
'select',
'six',
'six.moves',
'socket',
'subprocess',
'sys',
Expand Down
17 changes: 17 additions & 0 deletions pyupgrade/_plugins/six_simple.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,16 @@ def visit_Attribute(

func = functools.partial(replace_name, name=node.attr, new=new)
yield ast_to_offset(node), func
elif (
state.settings.min_version >= (3,) and
isinstance(node.value, ast.Attribute) and
isinstance(node.value.value, ast.Name) and
node.value.value.id == 'six' and
node.value.attr == 'moves' and
node.attr in {'xrange', 'range'}
):
func = functools.partial(replace_name, name=node.attr, new='range')
yield ast_to_offset(node), func


@register(ast.Name)
Expand Down Expand Up @@ -95,3 +105,10 @@ def visit_Name(

func = functools.partial(replace_name, name=node.id, new=new)
yield ast_to_offset(node), func
elif (
state.settings.min_version >= (3,) and
node.id in state.from_imports['six.moves'] and
node.id in {'xrange', 'range'}
):
func = functools.partial(replace_name, name=node.id, new='range')
yield ast_to_offset(node), func
27 changes: 27 additions & 0 deletions tests/features/six_simple_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,10 @@
'isinstance("foo", text_type)\n',
id='relative import might not be six',
),
pytest.param(
'foo.range(3)',
id='Range, but not from six.moves',
),
),
)
def test_six_simple_noop(s):
Expand Down Expand Up @@ -63,6 +67,29 @@ def test_six_simple_noop(s):
'from six import string_types\n'
'STRING_TYPES = (str,)\n',
),
pytest.param(
'six.moves.range(3)\n',
'range(3)\n',
id='six.moves.range',
),
pytest.param(
'six.moves.xrange(3)\n',
'range(3)\n',
id='six.moves.xrange',
),
pytest.param(
'from six.moves import xrange\n'
'xrange(3)\n',
'from six.moves import xrange\n'
'range(3)\n',
id='six.moves.xrange, from import',
),
),
)
def test_fix_six_simple(s, expected):
Expand Down

0 comments on commit 943891c

Please sign in to comment.