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 _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 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 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 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 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 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 _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 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=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 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 _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 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, match brackets 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)
comment_splicer.SpliceComments(tree)
if dumptree:
pytree_visitor.DumpPyTree(tree, target_stream=sys.stderr)
return pytree_unwrapper.UnwrapPyTree(tree)
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)
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:
Tuple of (reformatted_source, changed). reformatted_source conforms to the
desired formatting style. changed is True if the source changed.
"""
_CheckPythonVersion()
if not unformatted_source.endswith('\n'):
unformatted_source += '\n'
fst_newline = unformatted_source.find('\n')
indent_match = re.match('^(\s+)', unformatted_source[:fst_newline])
st = style.CreateStyleFromConfig(style_config)
if indent_match:
unformatted_source = dedent(unformatted_source)
original_indent = indent_match.group(0)
original_len = len(original_indent)
offset = original_len - original_len % st['INDENT_WIDTH']
probable_indent = original_indent[:offset]
st['COLUMN_LIMIT'] -= len(probable_indent)
style.SetGlobalStyle(st)
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 indent_match:
reformatted_source = indent(reformatted_source, probable_indent)
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 != ''
return reformatted_source, True
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.
"""
_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
