bpo-28837: Fix lib2to3 handling of map/zip/filter calls when followed with a 'trailer', e.g. zip()[x] (#24)

Author: stuarteberg

Lib/lib2to3/fixes/fix_filter.py

@@ -15,7 +15,10 @@
 
 # Local imports
 from .. import fixer_base
-from ..fixer_util import Name, Call, ListComp, in_special_context
+from ..pytree import Node
+from ..pygram import python_symbols as syms
+from ..fixer_util import Name, ArgList, ListComp, in_special_context
+
 
 class FixFilter(fixer_base.ConditionalFix):
  BM_compatible = True
@@ -34,16 +37,19 @@ class FixFilter(fixer_base.ConditionalFix):
  >
  ')'
  >
+ [extra_trailers=trailer*]
  >
  |
  power<
  'filter'
  trailer< '(' arglist< none='None' ',' seq=any > ')' >
+ [extra_trailers=trailer*]
  >
  |
  power<
  'filter'
  args=trailer< '(' [any] ')' >
+ [extra_trailers=trailer*]
  >
  """
 
@@ -53,23 +59,32 @@ def transform(self, node, results):
  if self.should_skip(node):
  return
 
+ trailers = []
+ if 'extra_trailers' in results:
+ for t in results['extra_trailers']:
+ trailers.append(t.clone())
+
  if "filter_lambda" in results:
  new = ListComp(results.get("fp").clone(),
  results.get("fp").clone(),
  results.get("it").clone(),
  results.get("xp").clone())
+ new = Node(syms.power, [new] + trailers, prefix="")
 
  elif "none" in results:
  new = ListComp(Name("_f"),
  Name("_f"),
  results["seq"].clone(),
  Name("_f"))
+ new = Node(syms.power, [new] + trailers, prefix="")
 
  else:
  if in_special_context(node):
  return None
- new = node.clone()
+
+ args = results['args'].clone()
+ new = Node(syms.power, [Name("filter"), args], prefix="")
+ new = Node(syms.power, [Name("list"), ArgList([new])] + trailers)
  new.prefix = ""
- new = Call(Name("list"), [new])
  new.prefix = node.prefix
  return new

Lib/lib2to3/fixes/fix_map.py

@@ -22,8 +22,10 @@
 # Local imports
 from ..pgen2 import token
 from .. import fixer_base
-from ..fixer_util import Name, Call, ListComp, in_special_context
+from ..fixer_util import Name, ArgList, Call, ListComp, in_special_context
 from ..pygram import python_symbols as syms
+from ..pytree import Node
+
 
 class FixMap(fixer_base.ConditionalFix):
  BM_compatible = True
@@ -32,6 +34,7 @@ class FixMap(fixer_base.ConditionalFix):
  map_none=power<
  'map'
  trailer< '(' arglist< 'None' ',' arg=any [','] > ')' >
+ [extra_trailers=trailer*]
  >
  |
  map_lambda=power<
@@ -47,10 +50,12 @@ class FixMap(fixer_base.ConditionalFix):
  >
  ')'
  >
+ [extra_trailers=trailer*]
  >
  |
  power<
- 'map' trailer< '(' [arglist=any] ')' >
+ 'map' args=trailer< '(' [any] ')' >
+ [extra_trailers=trailer*]
  >
  """
 
@@ -60,6 +65,11 @@ def transform(self, node, results):
  if self.should_skip(node):
  return
 
+ trailers = []
+ if 'extra_trailers' in results:
+ for t in results['extra_trailers']:
+ trailers.append(t.clone())
+
  if node.parent.type == syms.simple_stmt:
  self.warning(node, "You should use a for loop here")
  new = node.clone()
@@ -69,23 +79,32 @@ def transform(self, node, results):
  new = ListComp(results["xp"].clone(),
  results["fp"].clone(),
  results["it"].clone())
+ new = Node(syms.power, [new] + trailers, prefix="")
+
  else:
  if "map_none" in results:
  new = results["arg"].clone()
+ new.prefix = ""
  else:
- if "arglist" in results:
- args = results["arglist"]
- if args.type == syms.arglist and \
- args.children[0].type == token.NAME and \
- args.children[0].value == "None":
+ if "args" in results:
+ args = results["args"]
+ if args.type == syms.trailer and \
+ args.children[1].type == syms.arglist and \
+ args.children[1].children[0].type == token.NAME and \
+ args.children[1].children[0].value == "None":
  self.warning(node, "cannot convert map(None, ...) "
  "with multiple arguments because map() "
  "now truncates to the shortest sequence")
  return
+
+ new = Node(syms.power, [Name("map"), args.clone()])
+ new.prefix = ""
+
  if in_special_context(node):
  return None
- new = node.clone()
+
+ new = Node(syms.power, [Name("list"), ArgList([new])] + trailers)
  new.prefix = ""
- new = Call(Name("list"), [new])
+
  new.prefix = node.prefix
  return new

Lib/lib2to3/fixes/fix_zip.py

@@ -9,13 +9,16 @@
 
 # Local imports
 from .. import fixer_base
-from ..fixer_util import Name, Call, in_special_context
+from ..pytree import Node
+from ..pygram import python_symbols as syms
+from ..fixer_util import Name, ArgList, in_special_context
+
 
 class FixZip(fixer_base.ConditionalFix):
 
  BM_compatible = True
  PATTERN = """
- power< 'zip' args=trailer< '(' [any] ')' >
+ power< 'zip' args=trailer< '(' [any] ')' > [trailers=trailer*]
  >
  """
 
@@ -28,8 +31,16 @@ def transform(self, node, results):
  if in_special_context(node):
  return None
 
- new = node.clone()
- new.prefix = ""
- new = Call(Name("list"), [new])
+ args = results['args'].clone()
+ args.prefix = ""
+
+ trailers = []
+ if 'trailers' in results:
+ trailers = [n.clone() for n in results['trailers']]
+ for n in trailers:
+ n.prefix = ""
+
+ new = Node(syms.power, [Name("zip"), args], prefix="")
+ new = Node(syms.power, [Name("list"), ArgList([new])] + trailers)
  new.prefix = node.prefix
  return new

Lib/lib2to3/tests/test_fixers.py

@@ -2954,10 +2954,23 @@ def test_filter_basic(self):
  a = """x = [x for x in range(10) if x%2 == 0]"""
  self.check(b, a)
 
- # XXX This (rare) case is not supported
-## b = """x = filter(f, 'abc')[0]"""
-## a = """x = list(filter(f, 'abc'))[0]"""
-## self.check(b, a)
+ def test_filter_trailers(self):
+ b = """x = filter(None, 'abc')[0]"""
+ a = """x = [_f for _f in 'abc' if _f][0]"""
+ self.check(b, a)
+
+ b = """x = len(filter(f, 'abc')[0])"""
+ a = """x = len(list(filter(f, 'abc'))[0])"""
+ self.check(b, a)
+
+ b = """x = filter(lambda x: x%2 == 0, range(10))[0]"""
+ a = """x = [x for x in range(10) if x%2 == 0][0]"""
+ self.check(b, a)
+
+ # Note the parens around x
+ b = """x = filter(lambda (x): x%2 == 0, range(10))[0]"""
+ a = """x = [x for x in range(10) if x%2 == 0][0]"""
+ self.check(b, a)
 
  def test_filter_nochange(self):
  a = """b.join(filter(f, 'abc'))"""
@@ -3022,6 +3035,23 @@ def test_prefix_preservation(self):
  a = """x = list(map( f, 'abc' ))"""
  self.check(b, a)
 
+ def test_map_trailers(self):
+ b = """x = map(f, 'abc')[0]"""
+ a = """x = list(map(f, 'abc'))[0]"""
+ self.check(b, a)
+
+ b = """x = map(None, l)[0]"""
+ a = """x = list(l)[0]"""
+ self.check(b, a)
+
+ b = """x = map(lambda x:x, l)[0]"""
+ a = """x = [x for x in l][0]"""
+ self.check(b, a)
+
+ b = """x = map(f, 'abc')[0][1]"""
+ a = """x = list(map(f, 'abc'))[0][1]"""
+ self.check(b, a)
+
  def test_trailing_comment(self):
  b = """x = map(f, 'abc') # foo"""
  a = """x = list(map(f, 'abc')) # foo"""
@@ -3066,11 +3096,6 @@ def test_map_basic(self):
  """
  self.warns(b, a, "You should use a for loop here")
 
- # XXX This (rare) case is not supported
-## b = """x = map(f, 'abc')[0]"""
-## a = """x = list(map(f, 'abc'))[0]"""
-## self.check(b, a)
-
  def test_map_nochange(self):
  a = """b.join(map(f, 'abc'))"""
  self.unchanged(a)
@@ -3130,6 +3155,10 @@ def check(self, b, a):
  super(Test_zip, self).check(b, a)
 
  def test_zip_basic(self):
+ b = """x = zip()"""
+ a = """x = list(zip())"""
+ self.check(b, a)
+
  b = """x = zip(a, b, c)"""
  a = """x = list(zip(a, b, c))"""
  self.check(b, a)
@@ -3138,6 +3167,15 @@ def test_zip_basic(self):
  a = """x = len(list(zip(a, b)))"""
  self.check(b, a)
 
+ def test_zip_trailers(self):
+ b = """x = zip(a, b, c)[0]"""
+ a = """x = list(zip(a, b, c))[0]"""
+ self.check(b, a)
+
+ b = """x = zip(a, b, c)[0][1]"""
+ a = """x = list(zip(a, b, c))[0][1]"""
+ self.check(b, a)
+
  def test_zip_nochange(self):
  a = """b.join(zip(a, b))"""
  self.unchanged(a)

Misc/ACKS

@@ -128,6 +128,7 @@ Andrew Bennetts
 Andy Bensky
 Bennett Benson
 Ezra Berch
+Stuart Berg
 Michel Van den Bergh
 Julian Berman
 Brice Berna