def wrap_function_def(asttokens: ASTTokens,
node: ast.FunctionDef) -> WrappingSummary:
positions = node_start_positions(asttokens, node.args.args)
if node.args.vararg:
# Account for the * before the name
args_star = asttokens.prev_token(_first_token(node.args.vararg))
positions.append(Position(*args_star.start))
if node.args.kwonlyargs:
# Account for the unnamed *
if not node.args.vararg:
comma = asttokens.prev_token(_first_token(node.args.kwonlyargs[0]))
args_star = asttokens.prev_token(comma)
positions.append(Position(*args_star.start))
positions += node_start_positions(asttokens, node.args.kwonlyargs)
if node.args.kwarg:
# Account for the ** before the name
kwargs_stars = asttokens.prev_token(_first_token(node.args.kwarg))
positions.append(Position(*kwargs_stars.start))
summary = [(Position(pos.line, pos.col), MutationType.WRAP_INDENT)
for pos in positions]
close_paren = asttokens.next_token(_last_token(node.args))
args_end = Position(*close_paren.start)
if not (node.args.kwonlyargs or node.args.kwarg):
summary.append((args_end, MutationType.TRAILING_COMMA))
summary.append((args_end, MutationType.WRAP))
return summary
def node_start_position(asttokens: ASTTokens, node: ast.AST) -> Position:
first_token = _first_token(node)
if (isinstance(node, ast.GeneratorExp)
and generator_is_parenthesised(asttokens, node)):
first_token = asttokens.prev_token(first_token)
elif (isinstance(node, ast.BoolOp)
and bool_op_is_parenthesised(asttokens, node)):
first_token = asttokens.prev_token(first_token)
return Position(*first_token.start)
def generator_is_parenthesised(asttokens: ASTTokens,
node: ast.GeneratorExp) -> bool:
prev_token = asttokens.prev_token(_first_token(node))
next_token = asttokens.next_token(_last_token(node))
if prev_token.string == '(' and next_token.string == ')':
# These parens might be wrapping us
prev_prev_token = asttokens.prev_token(prev_token)
if prev_prev_token.string in ('(', ','):
return True
return False
def bool_op_is_parenthesised(asttokens: ASTTokens, node: ast.BoolOp) -> bool:
prev_token = asttokens.prev_token(_first_token(node))
next_token = asttokens.next_token(_last_token(node))
if prev_token.string == '(' and next_token.string == ')':
return True
return False
def wrap_if_exp(asttokens: ASTTokens, node: ast.IfExp) -> WrappingSummary:
if_token = asttokens.find_token(_first_token(node.test),
token.NAME,
'if',
reverse=True)
else_token = asttokens.find_token(_last_token(node.test), token.NAME,
'else')
summary = [(
Position(*_first_token(node).start),
MutationType.WRAP_INDENT,
), (
Position(*if_token.start),
MutationType.WRAP_INDENT,
), (
Position(*else_token.start),
MutationType.WRAP_INDENT,
), (
Position(*_last_token(node).end),
MutationType.WRAP,
)]
# Work out if we have parentheses already, if not we need to add some
if asttokens.prev_token(_first_token(node)).string != '(':
summary.insert(0, (
Position.from_node_start(node),
MutationType.OPEN_PAREN,
))
summary.append((
Position(*_last_token(node).end),
MutationType.CLOSE_PAREN,
))
return summary
def wrap_class_def(asttokens: ASTTokens,
node: ast.ClassDef) -> WrappingSummary:
if not node.bases and not node.keywords:
return []
named_args = node.keywords
kwargs = None
if named_args and named_args[-1].arg is None:
named_args = node.keywords[:-1]
kwargs = node.keywords[-1]
args = [*node.bases, *named_args]
summary = wrap_node_start_positions(asttokens, args)
if kwargs is not None:
kwargs_stars = asttokens.prev_token(_first_token(kwargs))
summary.append(
(Position(*kwargs_stars.start), MutationType.WRAP_INDENT))
summary.append(
(Position(*_last_token(kwargs).end), MutationType.TRAILING_COMMA))
summary.append((Position(*_last_token(kwargs).end), MutationType.WRAP))
else:
last_token_before_body = asttokens.next_token(_last_token(args[-1]))
summary.append((
Position(*last_token_before_body.start),
MutationType.TRAILING_COMMA,
))
summary.append((
Position(*last_token_before_body.start),
MutationType.WRAP,
))
return summary
def get_current_indent(asttokens: ASTTokens, node: ast.AST) -> int:
first_token = _first_token(node)
lineno = first_token.start[0]
next_tok = tok = first_token
while lineno == tok.start[0] and tok.type != token.INDENT:
next_tok = tok
tok = asttokens.prev_token(tok)
return next_tok.start[1] # type: ignore
def wrap_dict(asttokens: ASTTokens, node: ast.Dict) -> WrappingSummary:
positions = []
for key, value in zip(node.keys, node.values):
if key is not None:
positions.append(Position.from_node_start(key))
else:
kwargs_stars = asttokens.prev_token(_first_token(value))
positions.append(Position(*kwargs_stars.start))
summary = [(x, MutationType.WRAP_INDENT) for x in positions]
append_trailing_comma(asttokens, summary, node)
append_wrap_end(summary, node)
return summary
def wrap_generator_body(
asttokens: ASTTokens,
elt: ast.expr,
generators: List[ast.comprehension],
) -> WrappingSummary:
start_positions = [Position.from_node_start(elt)]
for generator in generators:
start_positions.append(Position.from_node_start(generator))
for compare in generator.ifs:
if_token = asttokens.prev_token(_first_token(compare))
assert if_token.string == 'if'
start_positions.append(Position(*if_token.start))
return [(x, MutationType.WRAP_INDENT) for x in start_positions]
def wrap_bool_op(asttokens: ASTTokens, node: ast.BoolOp) -> WrappingSummary:
summary = wrap_node_start_positions(asttokens, node.values)
summary.append((
Position(*_last_token(node).end),
MutationType.WRAP,
))
# Work out if we have parentheses already, if not we need to add some
if asttokens.prev_token(_first_token(node)).string != '(':
summary.insert(0, (
Position.from_node_start(node),
MutationType.OPEN_PAREN,
))
summary.append((
Position(*_last_token(node).end),
MutationType.CLOSE_PAREN,
))
return summary
def append_trailing_comma(
asttokens: ASTTokens,
summary: WrappingSummary,
node: ast.AST,
) -> WrappingSummary:
# Use the end position of the last content token, rather than the start
# position of the closing token. This ensures that we put the comma in the
# right place in constructs like:
#
# func(
# 'abcd', 'defg'
# )
#
# Where we want to put the comma immediately after 'defg' rather than just
# before the closing paren.
last_body_token = asttokens.prev_token(_last_token(node))
summary.append((
Position(*last_body_token.end),
MutationType.TRAILING_COMMA,
))
return summary
def wrap_call(asttokens: ASTTokens, node: ast.Call) -> WrappingSummary:
named_args = node.keywords
kwargs = None
if named_args and named_args[-1].arg is None:
named_args = node.keywords[:-1]
kwargs = node.keywords[-1]
if (len(node.args) == 1 and not named_args
and isinstance(node.args[0], ast.GeneratorExp)
and not generator_is_parenthesised(asttokens, node.args[0])):
generator_node = node.args[0] # type: ast.GeneratorExp
# The generator needs parentheses adding, as well as wrapping
summary = [(
Position.from_node_start(generator_node),
MutationType.WRAP_INDENT,
), (
Position.from_node_start(generator_node),
MutationType.OPEN_PAREN,
),
(
Position(*_last_token(generator_node).end),
MutationType.CLOSE_PAREN,
)]
else:
summary = wrap_node_start_positions(asttokens,
[*node.args, *named_args])
if kwargs is not None:
kwargs_stars = asttokens.prev_token(_first_token(kwargs))
summary.append(
(Position(*kwargs_stars.start), MutationType.WRAP_INDENT))
append_trailing_comma(asttokens, summary, node)
append_wrap_end(summary, node)
return summary
