fix_itertools.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. """ Fixer for itertools.(imap|ifilter|izip) --> (map|filter|zip) and
  2. itertools.ifilterfalse --> itertools.filterfalse (bugs 2360-2363)
  3. imports from itertools are fixed in fix_itertools_import.py
  4. If itertools is imported as something else (ie: import itertools as it;
  5. it.izip(spam, eggs)) method calls will not get fixed.
  6. """
  7. # Local imports
  8. from .. import fixer_base
  9. from ..fixer_util import Name
  10. class FixItertools(fixer_base.BaseFix):
  11. BM_compatible = True
  12. it_funcs = "('imap'|'ifilter'|'izip'|'izip_longest'|'ifilterfalse')"
  13. PATTERN = """
  14. power< it='itertools'
  15. trailer<
  16. dot='.' func=%(it_funcs)s > trailer< '(' [any] ')' > >
  17. |
  18. power< func=%(it_funcs)s trailer< '(' [any] ')' > >
  19. """ %(locals())
  20. # Needs to be run after fix_(map|zip|filter)
  21. run_order = 6
  22. def transform(self, node, results):
  23. prefix = None
  24. func = results['func'][0]
  25. if ('it' in results and
  26. func.value not in ('ifilterfalse', 'izip_longest')):
  27. dot, it = (results['dot'], results['it'])
  28. # Remove the 'itertools'
  29. prefix = it.prefix
  30. it.remove()
  31. # Replace the node which contains ('.', 'function') with the
  32. # function (to be consistent with the second part of the pattern)
  33. dot.remove()
  34. func.parent.replace(func)
  35. prefix = prefix or func.prefix
  36. func.replace(Name(func.value[1:], prefix=prefix))