Add fixer for itertools

docs/futurize.rst

@@ -103,6 +103,7 @@ The complete set of fixers applied by ``futurize --stage1`` is:
     lib2to3.fixes.fix_ws_comma
     lib2to3.fixes.fix_xreadlines
     libfuturize.fixes.fix_absolute_import
+    libfuturize.fixes.fix_itertools
     libfuturize.fixes.fix_next_call
     libfuturize.fixes.fix_print_with_import
     libfuturize.fixes.fix_raise

src/future/builtins/__init__.py

@@ -18,11 +18,14 @@
     import builtins
     bytes = builtins.bytes
     dict = builtins.dict
+    filter = builtins.filter
     int = builtins.int
     list = builtins.list
+    map = builtins.map
     object = builtins.object
     range = builtins.range
     str = builtins.str
+    zip = builtins.zip
     __all__ = []
 else:
     from future.types import (newbytes as bytes,
@@ -32,6 +35,9 @@
                               newobject as object,
                               newrange as range,
                               newstr as str)
+    from itertools import (ifilter as filter,
+                           imap as map,
+                           izip as zip)
 from future import utils
 
 

src/libfuturize/fixes/__init__.py

@@ -69,6 +69,7 @@
 
 libfuturize_fix_names_stage1 = set([
     'libfuturize.fixes.fix_absolute_import',
+    'libfuturize.fixes.fix_itertools',
     'libfuturize.fixes.fix_next_call',  # obj.next() -> next(obj). Unlike
                                         # lib2to3.fixes.fix_next, doesn't change
                                         # the ``next`` method to ``__next__``.

src/libfuturize/fixes/fix_itertools.py

@@ -0,0 +1,54 @@
+"""
+For the ``future`` package.
+
+Fixer for itertools methods that no longer deviate from builtins.
+
+This applies to imap, izip, and ifilter
+
+Adds this import line:
+
+    from builtins import filter, map, zip
+
+at the top.
+"""
+
+from lib2to3 import fixer_base
+
+from libfuturize.fixer_util import touch_import_top
+
+filter_expression = "name='ifilter'"
+map_expression = "name='imap'"
+zip_expression = "name='izip'"
+
+class FixFilter(fixer_base.BaseFix):
+
+    PATTERN = """
+              power<
+                 ({0}) trailer< '(' args=[any] ')' >
+              rest=any* >
+              """.format(filter_expression)
+
+    def transform(self, node, results):
+        touch_import_top(u'builtins', 'filter', node)
+
+class FixMap(fixer_base.BaseFix):
+
+    PATTERN = """
+              power<
+                 ({0}) trailer< '(' args=[any] ')' >
+              rest=any* >
+              """.format(map_expression)
+
+    def transform(self, node, results):
+        touch_import_top(u'builtins', 'map', node)
+
+class FixZip(fixer_base.BaseFix):
+
+    PATTERN = """
+              power<
+                 ({0}) trailer< '(' args=[any] ')' >
+              rest=any* >
+              """.format(zip_expression)
+
+    def transform(self, node, results):
+        touch_import_top(u'builtins', 'zip', node)

tests/test_future/test_futurize.py

@@ -453,6 +453,81 @@ def test_xrange(self):
         """
         self.convert_check(before, after, ignore_imports=False)
 
+    def test_filter(self):
+        """
+        Tests correct handling of itertools.ifilter
+        """
+        before = """
+        import itertools
+        itertools.ifilter(lambda x: x%2, [1, 2, 3, 4])
+        """
+        after = """
+        from builtins import filter
+        import itertools
+        filter(lambda x: x%2, [1, 2, 3, 4])
+        """
+        self.convert_check(before, after, ignore_imports=False)
+        before = """
+        from itertools import ifilter
+        ifilter(lambda x: x%2, [1, 2, 3, 4])
+        """
+        after = """
+        from builtins import filter
+
+        filter(lambda x: x%2, [1, 2, 3, 4])
+        """
+        self.convert_check(before, after, ignore_imports=False)
+
+    def test_map(self):
+        """
+        Tests correct handling of itertools.imap
+        """
+        before = """
+        import itertools
+        itertools.imap(pow, (2,3,10), (5,2,3))
+        """
+        after = """
+        from builtins import map
+        import itertools
+        map(pow, (2,3,10), (5,2,3))
+        """
+        self.convert_check(before, after, ignore_imports=False)
+        before = """
+        from itertools import imap
+        imap(pow, (2,3,10), (5,2,3))
+        """
+        after = """
+        from builtins import map
+
+        map(pow, (2,3,10), (5,2,3))
+        """
+        self.convert_check(before, after, ignore_imports=False)
+
+    def test_zip(self):
+        """
+        Tests correct handling of itertools.izip
+        """
+        before = """
+        import itertools
+        itertools.izip('ABCD', 'xy')
+        """
+        after = """
+        from builtins import zip
+        import itertools
+        zip('ABCD', 'xy')
+        """
+        self.convert_check(before, after, ignore_imports=False)
+        before = """
+        from itertools import izip
+        izip('ABCD', 'xy')
+        """
+        after = """
+        from builtins import zip
+
+        zip('ABCD', 'xy')
+        """
+        self.convert_check(before, after, ignore_imports=False)
+
     def test_source_coding_utf8(self):
         """
         Tests to ensure that the source coding line is not corrupted or