diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/lib2to3/fixes/fix_itertools.py | 5 | ||||
| -rw-r--r-- | Lib/lib2to3/fixes/fix_itertools_imports.py | 5 | ||||
| -rw-r--r-- | Lib/lib2to3/tests/test_fixers.py | 38 |
3 files changed, 32 insertions, 16 deletions
diff --git a/Lib/lib2to3/fixes/fix_itertools.py b/Lib/lib2to3/fixes/fix_itertools.py index 80790bf..63346b9 100644 --- a/Lib/lib2to3/fixes/fix_itertools.py +++ b/Lib/lib2to3/fixes/fix_itertools.py @@ -13,7 +13,7 @@ from ..fixer_util import Name class FixItertools(fixer_base.BaseFix): BM_compatible = True - it_funcs = "('imap'|'ifilter'|'izip'|'ifilterfalse')" + it_funcs = "('imap'|'ifilter'|'izip'|'izip_longest'|'ifilterfalse')" PATTERN = """ power< it='itertools' trailer< @@ -28,7 +28,8 @@ class FixItertools(fixer_base.BaseFix): def transform(self, node, results): prefix = None func = results['func'][0] - if 'it' in results and func.value != 'ifilterfalse': + if ('it' in results and + func.value not in ('ifilterfalse', 'izip_longest')): dot, it = (results['dot'], results['it']) # Remove the 'itertools' prefix = it.prefix diff --git a/Lib/lib2to3/fixes/fix_itertools_imports.py b/Lib/lib2to3/fixes/fix_itertools_imports.py index be8b90c..0ddbc7b 100644 --- a/Lib/lib2to3/fixes/fix_itertools_imports.py +++ b/Lib/lib2to3/fixes/fix_itertools_imports.py @@ -31,9 +31,10 @@ class FixItertoolsImports(fixer_base.BaseFix): if member_name in ('imap', 'izip', 'ifilter'): child.value = None child.remove() - elif member_name == 'ifilterfalse': + elif member_name in ('ifilterfalse', 'izip_longest'): node.changed() - name_node.value = 'filterfalse' + name_node.value = ('filterfalse' if member_name[1] == 'f' + else 'zip_longest') # Make sure the import statement is still sane children = imports.children[:] or [imports] diff --git a/Lib/lib2to3/tests/test_fixers.py b/Lib/lib2to3/tests/test_fixers.py index 0e11f38..43184e1 100644 --- a/Lib/lib2to3/tests/test_fixers.py +++ b/Lib/lib2to3/tests/test_fixers.py @@ -3623,16 +3623,24 @@ class Test_itertools(FixerTestCase): a = """%s(f, a)""" self.checkall(b, a) - def test_2(self): + def test_qualified(self): b = """itertools.ifilterfalse(a, b)""" a = """itertools.filterfalse(a, b)""" self.check(b, a) - def test_4(self): + b = """itertools.izip_longest(a, b)""" + a = """itertools.zip_longest(a, b)""" + self.check(b, a) + + def test_2(self): b = """ifilterfalse(a, b)""" a = """filterfalse(a, b)""" self.check(b, a) + b = """izip_longest(a, b)""" + a = """zip_longest(a, b)""" + self.check(b, a) + def test_space_1(self): b = """ %s(f, a)""" a = """ %s(f, a)""" @@ -3643,9 +3651,14 @@ class Test_itertools(FixerTestCase): a = """ itertools.filterfalse(a, b)""" self.check(b, a) + b = """ itertools.izip_longest(a, b)""" + a = """ itertools.zip_longest(a, b)""" + self.check(b, a) + def test_run_order(self): self.assert_runs_after('map', 'zip', 'filter') + class Test_itertools_imports(FixerTestCase): fixer = 'itertools_imports' @@ -3696,18 +3709,19 @@ class Test_itertools_imports(FixerTestCase): s = "from itertools import bar as bang" self.unchanged(s) - def test_ifilter(self): - b = "from itertools import ifilterfalse" - a = "from itertools import filterfalse" - self.check(b, a) + def test_ifilter_and_zip_longest(self): + for name in "filterfalse", "zip_longest": + b = "from itertools import i%s" % (name,) + a = "from itertools import %s" % (name,) + self.check(b, a) - b = "from itertools import imap, ifilterfalse, foo" - a = "from itertools import filterfalse, foo" - self.check(b, a) + b = "from itertools import imap, i%s, foo" % (name,) + a = "from itertools import %s, foo" % (name,) + self.check(b, a) - b = "from itertools import bar, ifilterfalse, foo" - a = "from itertools import bar, filterfalse, foo" - self.check(b, a) + b = "from itertools import bar, i%s, foo" % (name,) + a = "from itertools import bar, %s, foo" % (name,) + self.check(b, a) def test_import_star(self): s = "from itertools import *" |