fix_zip.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. """
  2. Fixer that changes zip(seq0, seq1, ...) into list(zip(seq0, seq1, ...)
  3. unless there exists a 'from future_builtins import zip' statement in the
  4. top-level namespace.
  5. We avoid the transformation if the zip() call is directly contained in
  6. iter(<>), list(<>), tuple(<>), sorted(<>), ...join(<>), or for V in <>:.
  7. """
  8. # Local imports
  9. from .. import fixer_base
  10. from ..pytree import Node
  11. from ..pygram import python_symbols as syms
  12. from ..fixer_util import Name, ArgList, in_special_context
  13. class FixZip(fixer_base.ConditionalFix):
  14. BM_compatible = True
  15. PATTERN = """
  16. power< 'zip' args=trailer< '(' [any] ')' > [trailers=trailer*]
  17. >
  18. """
  19. skip_on = "future_builtins.zip"
  20. def transform(self, node, results):
  21. if self.should_skip(node):
  22. return
  23. if in_special_context(node):
  24. return None
  25. args = results['args'].clone()
  26. args.prefix = ""
  27. trailers = []
  28. if 'trailers' in results:
  29. trailers = [n.clone() for n in results['trailers']]
  30. for n in trailers:
  31. n.prefix = ""
  32. new = Node(syms.power, [Name("zip"), args], prefix="")
  33. new = Node(syms.power, [Name("list"), ArgList([new])] + trailers)
  34. new.prefix = node.prefix
  35. return new