self.maybe_newline()
self.write(" " * self._indent + text)
- def write(self, text):
- """Append a piece of text"""
- self._source.append(text)
+ def write(self, *text):
+ """Add new source parts"""
+ self._source.extend(text)
@contextmanager
def buffered(self, buffer = None):
def visit_Lambda(self, node):
with self.require_parens(_Precedence.TEST, node):
- self.write("lambda ")
- self.traverse(node.args)
+ self.write("lambda")
+ with self.buffered() as buffer:
+ self.traverse(node.args)
+ if buffer:
+ self.write(" ", *buffer)
self.write(": ")
self.set_precedence(_Precedence.TEST, node.body)
self.traverse(node.body)
self.check_src_roundtrip("a[1, 2]")
self.check_src_roundtrip("a[(1, *a)]")
+ def test_lambda_parameters(self):
+ self.check_src_roundtrip("lambda: something")
+ self.check_src_roundtrip("four = lambda: 2 + 2")
+ self.check_src_roundtrip("lambda x: x * 2")
+ self.check_src_roundtrip("square = lambda n: n ** 2")
+ self.check_src_roundtrip("lambda x, y: x + y")
+ self.check_src_roundtrip("add = lambda x, y: x + y")
+ self.check_src_roundtrip("lambda x, y, /, z, q, *, u: None")
+ self.check_src_roundtrip("lambda x, *y, **z: None")
+
+
class DirectoryTestCase(ASTTestCase):
"""Test roundtrip behaviour on all files in Lib and Lib/test."""