Repository: asottile/pyupgrade

Issue:
next(six.itervalues(d)) is rewritten to next(d.values())
```console
$ cat example.py
import six


print(next(six.itervalues({1: 2})))
$ python example.py
2
```
```console
$ pyupgrade --py36-plus example.py
Rewriting example.py
```
```console
$ cat example.py
import six


print(next({1: 2}.values()))
$ python example.py
Traceback (most recent call last):
  File "example.py", line 4, in <module>
    print(next({1: 2}.values()))
TypeError: 'dict_values' object is not an iterator
```


Comments:

Patch:
diff --git a/pyupgrade.py b/pyupgrade.py
index d660a6d..fd897bc 100644
--- a/pyupgrade.py
+++ b/pyupgrade.py
@@ -1184,6 +1184,7 @@ class FindPy3Plus(ast.NodeVisitor):
         self.six_add_metaclass = set()  # type: Set[ast.ClassDef]
         self.six_b = set()  # type: Set[ast.Call]
         self.six_calls = {}  # type: Dict[Offset, ast.Call]
+        self.six_iter = {}  # type: Dict[Offset, ast.Call]
         self._previous_node = None  # type: Optional[ast.AST]
         self.six_raises = {}  # type: Dict[Offset, ast.Call]
         self.six_remove_decorators = set()  # type: Set[Offset]
@@ -1384,6 +1385,18 @@ class FindPy3Plus(ast.NodeVisitor):
             self.six_b.add(_ast_to_offset(node))
         elif self._is_six(node.func, SIX_CALLS) and not _starargs(node):
             self.six_calls[_ast_to_offset(node)] = node
+        elif (
+                isinstance(node.func, ast.Name) and
+                node.func.id == 'next' and
+                not _starargs(node) and
+                isinstance(node.args[0], ast.Call) and
+                self._is_six(
+                    node.args[0].func,
+                    ('iteritems', 'iterkeys', 'itervalues'),
+                ) and
+                not _starargs(node.args[0])
+        ):
+            self.six_iter[_ast_to_offset(node.args[0])] = node.args[0]
         elif (
                 isinstance(self._previous_node, ast.Expr) and
                 self._is_six(node.func, SIX_RAISES) and
@@ -1799,6 +1812,7 @@ def _fix_py3_plus(contents_text):  # type: (str) -> str
             visitor.six_add_metaclass,
             visitor.six_b,
             visitor.six_calls,
+            visitor.six_iter,
             visitor.six_raises,
             visitor.six_remove_decorators,
             visitor.six_simple,
@@ -1865,6 +1879,13 @@ def _fix_py3_plus(contents_text):  # type: (str) -> str
             ):
                 func_args, end = _parse_call_args(tokens, j)
                 _replace_call(tokens, i, end, func_args, SIX_B_TMPL)
+        elif token.offset in visitor.six_iter:
+            j = _find_open_paren(tokens, i)
+            func_args, end = _parse_call_args(tokens, j)
+            call = visitor.six_iter[token.offset]
+            assert isinstance(call.func, (ast.Name, ast.Attribute))
+            template = 'iter({})'.format(_get_tmpl(SIX_CALLS, call.func))
+            _replace_call(tokens, i, end, func_args, template)
         elif token.offset in visitor.six_calls:
             j = _find_open_paren(tokens, i)
             func_args, end = _parse_call_args(tokens, j)


Test patch:
diff --git a/tests/six_test.py b/tests/six_test.py
index a139e7f..ae53419 100644
--- a/tests/six_test.py
+++ b/tests/six_test.py
@@ -337,6 +337,16 @@ def test_fix_six_noop(s):
 
             id='add_metaclass, indented',
         ),
+        pytest.param(
+            'print(six.itervalues({1:2}))\n',
+            'print({1:2}.values())\n',
+            id='six.itervalues',
+        ),
+        pytest.param(
+            'print(next(six.itervalues({1:2})))\n',
+            'print(next(iter({1:2}.values())))\n',
+            id='six.itervalues inside next(...)',
+        ),
     ),
 )
 def test_fix_six(s, expected):