Merge pull request #494 from MarcoGorelli/six-range

rewrite six.moves.range

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

pyupgrade/_data.py

@@ -48,6 +48,7 @@ class State(NamedTuple):
     'mmap',
     'select',
     'six',
+    'six.moves',
     'socket',
     'subprocess',
     'sys',

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

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):