def testParseCodeToTree(self):
# Since ParseCodeToTree is a thin wrapper around underlying lib2to3
# functionality, only a sanity test here...
tree = pytree_utils.ParseCodeToTree('foo = 2\n')
self.assertEqual('file_input', pytree_utils.NodeName(tree))
self.assertEqual(2, len(tree.children))
self.assertEqual('simple_stmt', pytree_utils.NodeName(tree.children[0]))
def testPrintFunctionToTree(self):
tree = pytree_utils.ParseCodeToTree(
'print("hello world", file=sys.stderr)\n')
self.assertEqual('file_input', pytree_utils.NodeName(tree))
self.assertEqual(2, len(tree.children))
self.assertEqual('simple_stmt',
pytree_utils.NodeName(tree.children[0]))
def _ParseAndUnwrap(code, dumptree=False):
"""Produces unwrapped 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 unwrapped lines.
"""
style.SetGlobalStyle(style.CreateGoogleStyle())
tree = pytree_utils.ParseCodeToTree(code)
comment_splicer.SpliceComments(tree)
subtype_assigner.AssignSubtypes(tree)
split_penalty.ComputeSplitPenalties(tree)
blank_line_calculator.CalculateBlankLines(tree)
if dumptree:
pytree_visitor.DumpPyTree(tree, target_stream=sys.stderr)
uwlines = pytree_unwrapper.UnwrapPyTree(tree)
for uwl in uwlines:
uwl.CalculateFormattingInformation()
return uwlines
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 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 fix_formatting(filesource):
if not ('{' in filesource and ('[' in filesource or '(' in filesource)):
return filesource
tree = pytree_utils.ParseCodeToTree(filesource)
for node in tree.children:
walk_tree(node, node.children)
return six.text_type(tree)
def FormatCode(unformatted_source,
filename='<unknown>',
style_config=None,
lines=None,
print_diff=False,
verify=False):
"""Format a string of Python code.
This provides an alternative entry point to YAPF.
Arguments:
unformatted_source: (unicode) The code to format.
filename: (unicode) The name of the file being reformatted.
remaining arguments: see comment at the top of this module.
Returns:
Tuple of (reformatted_source, changed). reformatted_source conforms to the
desired formatting style. changed is True if the source changed.
"""
_CheckPythonVersion()
style.SetGlobalStyle(style.CreateStyleFromConfig(style_config))
if not unformatted_source.endswith('\n'):
unformatted_source += '\n'
try:
tree = pytree_utils.ParseCodeToTree(unformatted_source)
except parse.ParseError as e:
e.msg = filename + ': ' + e.msg
raise
# 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)
uwlines = pytree_unwrapper.UnwrapPyTree(tree)
for uwl in uwlines:
uwl.CalculateFormattingInformation()
lines = _LineRangesToSet(lines)
_MarkLinesToFormat(uwlines, lines)
reformatted_source = reformatter.Reformat(_SplitSemicolons(uwlines),
verify, lines)
if unformatted_source == reformatted_source:
return '' if print_diff else reformatted_source, False
code_diff = _GetUnifiedDiff(unformatted_source,
reformatted_source,
filename=filename)
if print_diff:
return code_diff, code_diff.strip() != '' # pylint: disable=g-explicit-bool-comparison
return reformatted_source, True
def testDumpPyTree(self):
# Similar sanity checking for the convenience wrapper DumpPyTree
tree = pytree_utils.ParseCodeToTree(_VISITOR_TEST_SIMPLE_CODE)
stream = py3compat.StringIO()
pytree_visitor.DumpPyTree(tree, target_stream=stream)
dump_output = stream.getvalue()
self.assertIn('file_input [3 children]', dump_output)
self.assertIn("NAME(Leaf(1, 'foo'))", dump_output)
self.assertIn("EQUAL(Leaf(22, '='))", dump_output)
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 testDumper(self):
# PyTreeDumper is mainly a debugging utility, so only do basic sanity
# checking.
tree = pytree_utils.ParseCodeToTree(_VISITOR_TEST_SIMPLE_CODE)
stream = py3compat.StringIO()
pytree_visitor.PyTreeDumper(target_stream=stream).Visit(tree)
dump_output = stream.getvalue()
self.assertIn('file_input [3 children]', dump_output)
self.assertIn("NAME(Leaf(1, 'foo'))", dump_output)
self.assertIn("EQUAL(Leaf(22, '='))", dump_output)
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 FormatCode(unformatted_source,
filename='<unknown>',
style_config=None,
lines=None,
print_diff=False):
"""Format a string of Python code.
This provides an alternative entry point to YAPF.
Arguments:
unformatted_source: (unicode) The code to format.
filename: (unicode) The name of the file being reformatted.
style_config, lines, print_diff: see comment at the top of this module.
Returns:
The code reformatted to conform to the desired formatting style.
"""
style.SetGlobalStyle(style.CreateStyleFromConfig(style_config))
tree = pytree_utils.ParseCodeToTree(unformatted_source.rstrip() + '\n')
# Run passes on the tree, modifying it in place.
comment_splicer.SpliceComments(tree)
subtype_assigner.AssignSubtypes(tree)
split_penalty.ComputeSplitPenalties(tree)
blank_line_calculator.CalculateBlankLines(tree)
uwlines = pytree_unwrapper.UnwrapPyTree(tree)
if not uwlines:
return ''
for uwl in uwlines:
uwl.CalculateFormattingInformation()
if lines is not None:
reformatted_source = _FormatLineSnippets(unformatted_source, uwlines, lines)
else:
lines = _LinesToFormat(uwlines)
if lines:
reformatted_source = _FormatLineSnippets(unformatted_source, uwlines,
lines)
else:
reformatted_source = reformatter.Reformat(uwlines)
if unformatted_source == reformatted_source:
return '' if print_diff else reformatted_source
code_diff = _GetUnifiedDiff(unformatted_source, reformatted_source,
filename=filename)
if print_diff:
return code_diff
return reformatted_source
def testCollectAllNodeNamesSimpleCode(self):
tree = pytree_utils.ParseCodeToTree(_VISITOR_TEST_SIMPLE_CODE)
collector = _NodeNameCollector()
collector.Visit(tree)
expected_names = [
'file_input',
'simple_stmt', 'expr_stmt', 'NAME', 'EQUAL', 'NAME', 'NEWLINE',
'simple_stmt', 'expr_stmt', 'NAME', 'EQUAL', 'NAME', 'NEWLINE',
'ENDMARKER'] # yapf: disable
self.assertEqual(expected_names, collector.all_node_names)
expected_name_node_values = ['foo', 'bar', 'baz', 'x']
self.assertEqual(expected_name_node_values, collector.name_node_values)
def _ParseAndUnwrap(self, code):
tree = pytree_utils.ParseCodeToTree(code)
comment_splicer.SpliceComments(tree)
subtype_assigner.AssignSubtypes(tree)
split_penalty.ComputeSplitPenalties(tree)
uwlines = pytree_unwrapper.UnwrapPyTree(tree)
for i, uwline in enumerate(uwlines):
uwlines[i] = unwrapped_line.UnwrappedLine(uwline.depth, [
ft for ft in uwline.tokens
if ft.name not in pytree_utils.NONSEMANTIC_TOKENS
])
return uwlines
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 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 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 FormatCode(unformatted_source,
filename='<unknown>',
style_config=None,
lines=None,
print_diff=False,
verify=False):
"""Format a string of Python code.
This provides an alternative entry point to YAPF.
Arguments:
unformatted_source: (unicode) The code to format.
filename: (unicode) The name of the file being reformatted.
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.
print_diff: (bool) Instead of returning the reformatted source, return a
diff that turns the formatted source into reformatter source.
verify: (bool) True if reformatted code should be verified for syntax.
Returns:
Tuple of (reformatted_source, changed). reformatted_source conforms to the
desired formatting style. changed is True if the source changed.
"""
try:
tree = pytree_utils.ParseCodeToTree(unformatted_source)
except Exception as e:
e.filename = filename
raise errors.YapfError(errors.FormatErrorMsg(e))
reformatted_source = FormatTree(tree,
style_config=style_config,
lines=lines,
verify=verify)
if unformatted_source == reformatted_source:
return '' if print_diff else reformatted_source, False
code_diff = _GetUnifiedDiff(unformatted_source,
reformatted_source,
filename=filename)
if print_diff:
return code_diff, code_diff.strip() != '' # pylint: disable=g-explicit-bool-comparison # noqa
return reformatted_source, True
def FormatCode(unformatted_source,
filename='<unknown>',
style_config=None,
lines=None,
print_diff=False,
verify=True):
"""Format a string of Python code.
This provides an alternative entry point to YAPF.
Arguments:
unformatted_source: (unicode) The code to format.
filename: (unicode) The name of the file being reformatted.
remaining arguments: see comment at the top of this module.
Returns:
The code reformatted to conform to the desired formatting style.
"""
_CheckPythonVersion()
style.SetGlobalStyle(style.CreateStyleFromConfig(style_config))
if not unformatted_source.endswith('\n'):
unformatted_source += '\n'
tree = pytree_utils.ParseCodeToTree(unformatted_source)
# Run passes on the tree, modifying it in place.
comment_splicer.SpliceComments(tree)
continuation_splicer.SpliceContinuations(tree)
subtype_assigner.AssignSubtypes(tree)
split_penalty.ComputeSplitPenalties(tree)
blank_line_calculator.CalculateBlankLines(tree)
uwlines = pytree_unwrapper.UnwrapPyTree(tree)
for uwl in uwlines:
uwl.CalculateFormattingInformation()
_MarkLinesToFormat(uwlines, lines)
reformatted_source = reformatter.Reformat(uwlines, verify)
if unformatted_source == reformatted_source:
return '' if print_diff else reformatted_source
code_diff = _GetUnifiedDiff(unformatted_source,
reformatted_source,
filename=filename)
if print_diff:
return code_diff
return reformatted_source
def testCollectAllNodeNamesNestedCode(self):
tree = pytree_utils.ParseCodeToTree(_VISITOR_TEST_NESTED_CODE)
collector = _NodeNameCollector()
collector.Visit(tree)
expected_names = [
'file_input',
'if_stmt', 'NAME', 'NAME', 'COLON',
'suite', 'NEWLINE',
'INDENT', 'if_stmt', 'NAME', 'NAME', 'COLON', 'suite', 'NEWLINE',
'INDENT', 'simple_stmt', 'return_stmt', 'NAME', 'NAME', 'NEWLINE',
'DEDENT', 'DEDENT', 'ENDMARKER'] # yapf: disable
self.assertEqual(expected_names, collector.all_node_names)
expected_name_node_values = ['if', 'x', 'if', 'y', 'return', 'z']
self.assertEqual(expected_name_node_values, collector.name_node_values)
def _ParseAndUnwrap(self, code):
"""Produces unwrapped 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
Returns:
List of unwrapped lines.
"""
tree = pytree_utils.ParseCodeToTree(code)
comment_splicer.SpliceComments(tree)
return pytree_unwrapper.UnwrapPyTree(tree)
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 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 _ParseAndComputePenalties(self, code, dumptree=False):
"""Parses the code and computes split penalties.
Arguments:
code: code to parse as a string
dumptree: if True, the parsed pytree (after penalty assignment) is dumped
to stderr. Useful for debugging.
Returns:
Parse tree.
"""
tree = pytree_utils.ParseCodeToTree(code)
split_penalty.ComputeSplitPenalties(tree)
if dumptree:
pytree_visitor.DumpPyTree(tree, target_stream=sys.stderr)
return tree
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 _ParseAndUnwrap(self, code, dumptree=False):
"""Produces unwrapped lines from the given code.
Parses the code into a tree, assigns subtypes 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 unwrapped lines.
"""
tree = pytree_utils.ParseCodeToTree(code)
subtype_assigner.AssignSubtypes(tree)
if dumptree:
pytree_visitor.DumpPyTree(tree, target_stream=sys.stderr)
return pytree_unwrapper.UnwrapPyTree(tree)