def testCommentBeforeDedentThreeLevel(self): code = textwrap.dedent(r''' if foo: if bar: z = 1 # comment 2 # comment 1 # comment 0 j = 2 ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) # comment 0 should go under the tree root self._AssertNodeIsComment(tree.children[1], '# comment 0') # comment 1 is in the first if_suite, right before the DEDENT if_suite_1 = self._FindNthChildNamed(tree, 'suite', n=1) self._AssertNodeIsComment(if_suite_1.children[-2], '# comment 1') self._AssertNodeType('DEDENT', if_suite_1.children[-1]) # comment 2 is in if_suite nested under the first if suite, # right before the DEDENT if_suite_2 = self._FindNthChildNamed(tree, 'suite', n=2) self._AssertNodeIsComment(if_suite_2.children[-2], '# comment 2') self._AssertNodeType('DEDENT', if_suite_2.children[-1])
def ParseAndUnwrap(code, dumptree=False): """Produces logical lines from the given code. Parses the code into a tree, performs comment splicing and runs the unwrapper. Arguments: code: code to parse as a string dumptree: if True, the parsed pytree (after comment splicing) is dumped to stderr. Useful for debugging. Returns: List of logical lines. """ tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) continuation_splicer.SpliceContinuations(tree) subtype_assigner.AssignSubtypes(tree) identify_container.IdentifyContainers(tree) split_penalty.ComputeSplitPenalties(tree) blank_line_calculator.CalculateBlankLines(tree) if dumptree: pytree_visitor.DumpPyTree(tree, target_stream=sys.stderr) llines = pytree_unwrapper.UnwrapPyTree(tree) for lline in llines: lline.CalculateFormattingInformation() return llines
def testExprComments(self): code = textwrap.dedent(r''' foo( # Request fractions of an hour. 948.0/3600, 20) ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) trailer = self._FindNthChildNamed(tree, 'trailer', 1) comment = trailer.children[1] self._AssertNodeIsComment(comment, '# Request fractions of an hour.')
def testSimpleInline(self): code = 'foo = 1 # and a comment\n' tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) expr = tree.children[0].children[0] # Check that the expected node is still expr_stmt, but now it has 4 children # (before comment splicing it had 3), the last child being the comment. self._AssertNodeType('expr_stmt', expr) self.assertEqual(4, len(expr.children)) comment_node = expr.children[3] self._AssertNodeIsComment(comment_node, '# and a comment')
def testTwoLineComment(self): code = textwrap.dedent(r''' foo = 1 # first comment # second comment bar = 2 ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) # This is similar to the single-line standalone comment. self.assertEqual(4, len(tree.children)) self._AssertNodeIsComment(tree.children[1])
def testSimpleSeparateLine(self): code = textwrap.dedent(r''' foo = 1 # first comment bar = 2 ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) # The comment should've been added to the root's children (now 4, including # the ENDMARKER in the end. self.assertEqual(4, len(tree.children)) comment_node = tree.children[1] self._AssertNodeIsComment(comment_node)
def testCommentBeforeDedent(self): code = textwrap.dedent(r''' if bar: z = 1 # a comment j = 2 ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) # The comment should go under the tree root, not under the 'if'. self._AssertNodeIsComment(tree.children[1]) if_suite = tree.children[0].children[3] self._AssertNodeType('DEDENT', if_suite.children[-1])
def testCommentIsLastChildInCompound(self): code = textwrap.dedent(r''' if x: foo = 1 # a comment ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) # Look into the suite node under the 'if'. We don't care about the DEDENT # leaf but the new COMMENT must be a child of the suite and after the # actual code leaf. if_suite = tree.children[0].children[3] self._AssertNodeType('DEDENT', if_suite.children[-1]) self._AssertNodeIsComment(if_suite.children[-2])
def testCommentIsFirstChildInCompound(self): code = textwrap.dedent(r''' if x: # a comment foo = 1 ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) # Look into the suite node under the 'if'. We don't care about the NEWLINE # leaf but the new COMMENT must be a child of the suite and before the # actual code leaf. if_suite = tree.children[0].children[3] self._AssertNodeType('NEWLINE', if_suite.children[0]) self._AssertNodeIsComment(if_suite.children[1])
def testCommentBeforeDedentTwoLevel(self): code = textwrap.dedent(r''' if foo: if bar: z = 1 # a comment y = 1 ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) if_suite = tree.children[0].children[3] # The comment is in the first if_suite, not the nested if under it. It's # right before the DEDENT self._AssertNodeIsComment(if_suite.children[-2]) self._AssertNodeType('DEDENT', if_suite.children[-1])
def testSeparateLineAfterInline(self): code = textwrap.dedent(r''' bar = 1 foo = 1 # inline comment # line comment ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) # The separate line comment should become a child of the root, while # the inline comment remains within its simple_node. sep_comment_node = tree.children[-2] self._AssertNodeIsComment(sep_comment_node, '# line comment') expr = tree.children[1].children[0] inline_comment_node = expr.children[-1] self._AssertNodeIsComment(inline_comment_node, '# inline comment')
def testMultipleCommentsInOneExpr(self): code = textwrap.dedent(r''' foo( # com 1 948.0/3600, # com 2 20 + 12 # com 3 ) ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) trailer = self._FindNthChildNamed(tree, 'trailer', 1) self._AssertNodeIsComment(trailer.children[1], '# com 1') arglist = self._FindNthChildNamed(tree, 'arglist', 1) self._AssertNodeIsComment(arglist.children[2], '# com 2') arith_expr = self._FindNthChildNamed(tree, 'arith_expr', 1) self._AssertNodeIsComment(arith_expr.children[-1], '# com 3')
def testMultipleBlockComments(self): code = textwrap.dedent(r''' # Block comment number 1 # Block comment number 2 def f(): pass ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) funcdef = tree.children[0] block_comment_1 = funcdef.children[0] self._AssertNodeIsComment(block_comment_1, '# Block comment number 1') block_comment_2 = funcdef.children[1] self._AssertNodeIsComment(block_comment_2, '# Block comment number 2')
def testCommentBeforeDedentTwoLevelImproperlyIndented(self): code = textwrap.dedent(r''' if foo: if bar: z = 1 # comment 2 y = 1 ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) # The comment here is indented by 3 spaces, which is unlike any of the # surrounding statement indentation levels. The splicer attaches it to the # "closest" parent with smaller indentation. if_suite = tree.children[0].children[3] # The comment is in the first if_suite, not the nested if under it. It's # right before the DEDENT self._AssertNodeIsComment(if_suite.children[-2]) self._AssertNodeType('DEDENT', if_suite.children[-1])
def testCommentsInClass(self): code = textwrap.dedent(r''' class Foo: """docstring abc...""" # top-level comment def foo(): pass # another comment ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) class_suite = tree.children[0].children[3] another_comment = class_suite.children[-2] self._AssertNodeIsComment(another_comment, '# another') # It's OK for the comment to be a child of funcdef, as long as it's # the first child and thus comes before the 'def'. funcdef = class_suite.children[3] toplevel_comment = funcdef.children[0] self._AssertNodeIsComment(toplevel_comment, '# top-level')
def testCommentsOnDedents(self): code = textwrap.dedent(r''' class Foo(object): # A comment for qux. def qux(self): pass # Interim comment. def mux(self): pass ''') tree = pytree_utils.ParseCodeToTree(code) comment_splicer.SpliceComments(tree) classdef = tree.children[0] class_suite = classdef.children[6] qux_comment = class_suite.children[1] self._AssertNodeIsComment(qux_comment, '# A comment for qux.') interim_comment = class_suite.children[4] self._AssertNodeIsComment(interim_comment, '# Interim comment.')
def FormatTree(tree, style_config=None, lines=None, verify=False): """Format a parsed lib2to3 pytree. This provides an alternative entry point to YAPF. Arguments: tree: (pytree.Node) The root of the pytree to format. style_config: (string) Either a style name or a path to a file that contains formatting style settings. If None is specified, use the default style as set in style.DEFAULT_STYLE_FACTORY lines: (list of tuples of integers) A list of tuples of lines, [start, end], that we want to format. The lines are 1-based indexed. It can be used by third-party code (e.g., IDEs) when reformatting a snippet of code rather than a whole file. verify: (bool) True if reformatted code should be verified for syntax. Returns: The source formatted according to the given formatting style. """ _CheckPythonVersion() style.SetGlobalStyle(style.CreateStyleFromConfig(style_config)) # Run passes on the tree, modifying it in place. comment_splicer.SpliceComments(tree) continuation_splicer.SpliceContinuations(tree) subtype_assigner.AssignSubtypes(tree) identify_container.IdentifyContainers(tree) split_penalty.ComputeSplitPenalties(tree) blank_line_calculator.CalculateBlankLines(tree) llines = pytree_unwrapper.UnwrapPyTree(tree) for lline in llines: lline.CalculateFormattingInformation() lines = _LineRangesToSet(lines) _MarkLinesToFormat(llines, lines) return reformatter.Reformat(_SplitSemicolons(llines), verify, lines)