Merge pull request #120 from takluyver/dict-methods-in-loops

Fix transformation of non-iter dict methods in for loops
PyCQA · Mar 21, 2015 · c350457 · c350457
2 parents 5ab6f01 + 8a1ef24
commit c350457
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/libmodernize/fixes/fix_dict_six.py b/libmodernize/fixes/fix_dict_six.py
@@ -24,3 +24,9 @@ def transform(self, node, results):
             return super(FixDictSix, self).transform(node, results)
         else:
             return self.transform_iter(method_name, node, results['head'])
+
+    def in_special_context(self, node, isiter):
+        # Redefined from parent class to make "for x in d.items()" count as
+        # in special context; 2to3 only counts for loops as special context
+        # for the iter* methods.
+        return super(FixDictSix, self).in_special_context(node, True)
diff --git a/tests/test_fix_dict_six.py b/tests/test_fix_dict_six.py
@@ -27,6 +27,14 @@
 list(x.{type}())
 """)
 
+DICT_IN_LOOP = ("""\
+for k in x.items():
+    pass
+""", """\
+for k in x.items():
+    pass
+""")
+
 
 def check_all_types(input, output):
     for type_ in TYPES:
@@ -40,3 +48,6 @@ def test_dict_view():
 
 def test_dict_plain():
     check_all_types(*DICT_PLAIN)
+
+def test_dict_in_loop():
+    check_on_input(*DICT_IN_LOOP)