diff --git a/README.md b/README.md index 7bba52da..94f27c87 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/pyupgrade/_data.py b/pyupgrade/_data.py index 6218282d..e82fa464 100644 --- a/pyupgrade/_data.py +++ b/pyupgrade/_data.py @@ -48,6 +48,7 @@ class State(NamedTuple): 'mmap', 'select', 'six', + 'six.moves', 'socket', 'subprocess', 'sys', diff --git a/pyupgrade/_plugins/six_simple.py b/pyupgrade/_plugins/six_simple.py index 1e08e0dd..afa820f8 100644 --- a/pyupgrade/_plugins/six_simple.py +++ b/pyupgrade/_plugins/six_simple.py @@ -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) @@ -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 diff --git a/tests/features/six_simple_test.py b/tests/features/six_simple_test.py index d78835e4..a57d33c6 100644 --- a/tests/features/six_simple_test.py +++ b/tests/features/six_simple_test.py @@ -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): @@ -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):