Merge pull request #181 from martindemello/six-dict

Fix for dict.viewitems(), dict.iteritems() etc in chained calls.

Author: takluyver

libmodernize/fixes/fix_dict_six.py

@@ -3,27 +3,46 @@
 
 # Local imports
 from lib2to3 import fixer_util
+from lib2to3 import pytree
 from lib2to3.fixes import fix_dict
 import libmodernize
 
 
 class FixDictSix(fix_dict.FixDict):
 
-    def transform_iter(self, method_name, node, base):
+    def transform_iter(self, node, results):
         """Call six.(iter|view)items() and friends."""
+        # Make sure six is imported.
         libmodernize.touch_import(None, u'six', node)
-        new_node = [n.clone() for n in base]
-        new_node[0].prefix = u''
+
+        # Copy of self.transform() from lib2to3.fix_dict with some changes to
+        # use the six.* methods.
+
+        head = results['head']
+        method = results['method'][0] # Extract node for method name
+        tail = results['tail']
+        syms = self.syms
+        method_name = method.value
         name = fixer_util.Name(u'six.' + method_name, prefix=node.prefix)
-        node.replace(fixer_util.Call(name, new_node))
+        assert method_name.startswith((u'iter', u'view')), repr(method)
+        assert method_name[4:] in (u'keys', u'items', u'values'), repr(method)
+        head = [n.clone() for n in head]
+        tail = [n.clone() for n in tail]
+        new = pytree.Node(syms.power, head)
+        new.prefix = u''
+        new = fixer_util.Call(name, [new])
+        if tail:
+            new = pytree.Node(syms.power, [new] + tail)
+        new.prefix = node.prefix
+        return new
 
     def transform(self, node, results):
         method = results['method'][0]
         method_name = method.value
         if method_name in ('keys', 'items', 'values'):
             return super(FixDictSix, self).transform(node, results)
         else:
-            return self.transform_iter(method_name, node, results['head'])
+            return self.transform_iter(node, results)
 
     def in_special_context(self, node, isiter):
         # Redefined from parent class to make "for x in d.items()" count as

tests/test_fix_dict_six.py

@@ -35,6 +35,34 @@
     pass
 """)
 
+DICT_ITER_IN_LOOP = ("""\
+for k in x.iter{type}():
+    pass
+""", """\
+from __future__ import absolute_import
+import six
+for k in six.iter{type}(x):
+    pass
+""")
+
+DICT_ITER_IN_LIST = ("""\
+for k in list(x.iter{type}()):
+    pass
+""", """\
+from __future__ import absolute_import
+import six
+for k in list(six.iter{type}(x)):
+    pass
+""")
+
+CHAINED_CALLS = ("""\
+(x + y).foo().iter{type}().bar()
+""", """\
+from __future__ import absolute_import
+import six
+six.iter{type}((x + y).foo()).bar()
+""")
+
 
 def check_all_types(input, output):
     for type_ in TYPES:
@@ -51,3 +79,12 @@ def test_dict_plain():
 
 def test_dict_in_loop():
     check_on_input(*DICT_IN_LOOP)
+
+def test_dict_iter_in_loop():
+    check_all_types(*DICT_ITER_IN_LOOP)
+
+def test_dict_iter_in_list():
+    check_all_types(*DICT_ITER_IN_LIST)
+
+def test_chained_calls():
+    check_all_types(*CHAINED_CALLS)