def add_provide_context_to_python_operator(node: LN, capture: Capture, filename: Filename) -> None:
fn_args = capture['function_arguments'][0]
fn_args.append_child(Comma())
provide_context_arg = KeywordArg(Name('provide_context'), Name('True'))
provide_context_arg.prefix = fn_args.children[0].prefix
fn_args.append_child(provide_context_arg)
def add_provide_context_to_python_operator(node: LN, capture: Capture, filename: Filename) -> None:
fn_args = capture['function_arguments'][0]
if len(fn_args.children) > 0 and (not isinstance(fn_args.children[-1], Leaf)
or fn_args.children[-1].type != token.COMMA):
fn_args.append_child(Comma())
provide_context_arg = KeywordArg(Name('provide_context'), Name('True'))
provide_context_arg.prefix = fn_args.children[0].prefix
fn_args.append_child(provide_context_arg)
def on_delete_modifier(node, capture, filename) -> None:
if not capture:
return
args: Node = capture.get("function_arguments")[0]
has_on_delete = False
for node in args.children:
if isinstance(node, Node) and node.type == 261:
if node.children[0].value == "on_delete":
has_on_delete = True
if not has_on_delete:
if args.children[-1].value != ",":
args.append_child(Comma())
args.append_child(
KeywordArg(Name("on_delete", prefix=" "), Name("models.CASCADE"))
)
def pytest_approx(node, capture, filename):
target_value = listify(capture['target_value'])[0].clone()
target_value.prefix = ''
abs_tolerance = capture['abs_tolerance'].clone()
abs_tolerance.prefix = ''
op_value = listify(capture['op'])[0].value
# Adds a 'import pytest' if there wasn't one already
touch_import(None, "pytest", node)
if op_value in ('<', '<='):
# as you'd expect in an assert statement
operator = Leaf(TOKEN.EQEQUAL, '==', prefix=' ')
else:
# probably in an if statement
operator = Leaf(TOKEN.NOTEQUAL, '!=', prefix=' ')
node.replace(
Node(
syms.comparison,
[
capture['lhs'].clone(),
operator,
Node(
syms.power,
[
kw('pytest'),
Node(
syms.trailer,
[Leaf(TOKEN.DOT, ".", prefix=''), kw('approx', prefix='')],
prefix='',
),
ArgList(
[
target_value,
Comma(),
KeywordArg(kw('abs'), abs_tolerance),
]
),
],
),
],
prefix=node.prefix,
)
)
def add_argument(filename, trailer_node, key, value):
"""
add "key=value" to arglist in trailer_node,
if arglist already contains key, reassign key to value
"""
if trailer_node.type != python_symbols.trailer and len(
trailer_node.children) != 3:
log_warning(
filename, trailer_node.get_lineno(),
"node type is not trailer or len(children) != 3. you may need to call norm_arglist first."
)
return
arglist_node = trailer_node.children[1]
if arglist_node.type != python_symbols.arglist:
log_warning(filename, trailer_node.get_lineno(),
"trailer_node.children[1] is not arglist.")
return
found_key = False
for node in arglist_node.children:
if node.type != python_symbols.argument:
continue
_key_node = node.children[0]
if _key_node.value == key:
found_key = True
_value_node = node.children[2]
if _value_node.value != value:
_value_node_copy = _value_node.clone()
_value_node_copy.type = token.NAME
_value_node_copy.value = value
_value_node.replace(_value_node_copy)
log_warning(
filename, arglist_node.get_lineno(),
'argument "{}" is reassigned to "{}"'.format(key, value))
break
if not found_key:
key_node = Name(key)
value_node = Name(value)
if arglist_node.children:
arglist_node.append_child(Comma())
key_node.prefix = " "
arg_node = KeywordArg(key_node, value_node)
arglist_node.append_child(arg_node)
log_warning(filename, arglist_node.get_lineno(),
'add argument "{}={}"'.format(key, value))
def assertalmostequal_to_assert(node, capture, arguments):
function_name = capture["function_name"].value
invert = function_name in INVERT_FUNCTIONS
function_name = SYNONYMS.get(function_name, function_name)
nargs = len(arguments)
if nargs < 2 or nargs > 5:
return None
def get_kwarg_value(index, name):
idx = 0
for arg in arguments:
if arg.type == syms.argument:
if arg.children[0].value == name:
return arg.children[2].clone()
else:
if idx == index:
return arg.clone()
idx += 1
return None
first = get_kwarg_value(0, "first")
second = get_kwarg_value(1, "second")
if first is None or second is None:
# Not sure what this is, leave it alone
return
places = get_kwarg_value(2, "places")
msg = get_kwarg_value(3, "msg")
delta = get_kwarg_value(4, "delta")
if delta is not None:
try:
abs_delta = float(delta.value)
except ValueError:
# this should be a number, give up.
return
else:
if places is None:
places = 7
else:
try:
places = int(places.value)
except (ValueError, AttributeError):
# this should be an int, give up.
return
abs_delta = "1e-%d" % places
arguments[1].prefix = ""
if invert:
op_token = Leaf(TOKEN.NOTEQUAL, "!=", prefix=" ")
else:
op_token = Leaf(TOKEN.EQEQUAL, "==", prefix=" ")
assert_test_nodes = [
Node(
syms.comparison,
[
arguments[0],
op_token,
Node(
syms.power,
Attr(keyword("pytest"), keyword("approx", prefix="")) + [
ArgList([
arguments[1],
Comma(),
KeywordArg(keyword("abs"),
Leaf(TOKEN.NUMBER, abs_delta)),
])
],
),
],
)
]
# Adds a 'import pytest' if there wasn't one already
touch_import(None, "pytest", node)
return Assert(assert_test_nodes,
msg.clone() if msg else None,
prefix=node.prefix)
def assertmethod_to_assert(node, capture, arguments):
"""
self.assertEqual(foo, bar, msg)
--> assert foo == bar, msg
self.assertTrue(foo, msg)
--> assert foo, msg
self.assertIsNotNone(foo, msg)
--> assert foo is not None, msg
.. etc
"""
function_name = capture["function_name"].value
invert = function_name in INVERT_FUNCTIONS
function_name = SYNONYMS.get(function_name, function_name)
num_arguments = ARGUMENTS[function_name]
if len(arguments) not in (num_arguments, num_arguments + 1):
# Not sure what this is. Leave it alone.
return None
if len(arguments) == num_arguments:
message = None
else:
message = arguments.pop()
if message.type == syms.argument:
# keyword argument (e.g. `msg=abc`)
message = message.children[2].clone()
if function_name == "assertIsInstance":
arguments[0].prefix = ""
assert_test_nodes = [
Call(keyword("isinstance"),
[arguments[0], Comma(), arguments[1]])
]
if invert:
assert_test_nodes.insert(0, keyword("not"))
elif function_name == "assertAlmostEqual":
arguments[1].prefix = ""
# TODO: insert the `import pytest` at the top of the file
if invert:
op_token = Leaf(TOKEN.NOTEQUAL, "!=", prefix=" ")
else:
op_token = Leaf(TOKEN.EQEQUAL, "==", prefix=" ")
assert_test_nodes = [
Node(
syms.comparison,
[
arguments[0],
op_token,
Node(
syms.power,
Attr(keyword("pytest"), keyword("approx", prefix="")) +
[
ArgList([
arguments[1],
Comma(),
KeywordArg(keyword("abs"),
Leaf(TOKEN.NUMBER, "1e-7")),
])
],
),
],
)
]
# Adds a 'import pytest' if there wasn't one already
touch_import(None, "pytest", node)
else:
op_tokens = OPERATORS[function_name]
if not isinstance(op_tokens, list):
op_tokens = [op_tokens]
op_tokens = [o.clone() for o in op_tokens]
if invert:
if not op_tokens:
op_tokens.append(keyword("not"))
elif op_tokens[0].type == TOKEN.NAME and op_tokens[0].value == "is":
op_tokens[0] = Node(
syms.comp_op,
[keyword("is"), keyword("not")], prefix=" ")
elif op_tokens[0].type == TOKEN.NAME and op_tokens[0].value == "in":
op_tokens[0] = Node(
syms.comp_op,
[keyword("not"), keyword("in")], prefix=" ")
elif op_tokens[0].type == TOKEN.EQEQUAL:
op_tokens[0] = Leaf(TOKEN.NOTEQUAL, "!=", prefix=" ")
if num_arguments == 2:
# a != b, etc.
assert_test_nodes = [arguments[0]] + op_tokens + [arguments[1]]
elif function_name == "assertTrue":
assert_test_nodes = op_tokens + [arguments[0]]
# not a
elif function_name == "assertIsNone":
# a is not None
assert_test_nodes = [arguments[0]] + op_tokens
return Assert(assert_test_nodes,
message.clone() if message else None,
prefix=node.prefix)
def default_transformer(node: LN, capture: Capture, filename: Filename):
fp = capture.get("function_parameters")
if fp and fp.children[1].type == SYMBOL.arglist:
arg_node = KeywordArg(Name("trans_arg"), Number("1"))
fp.children[1].append_child(Comma())
fp.children[1].append_child(arg_node)
def make_pytest_raises_blocks(node, capture, filename):
"""
Turns this:
try:
...
pytest.fail(...)
except:
pass
Into:
with pytest.raises(Exception):
...
Not only is this prettier, but the former is a bug since
pytest.fail() raises an exception.
"""
exc_class = capture.get('exc_class', None)
if exc_class:
exc_class = exc_class.clone()
exc_class.prefix = ''
raises_args = [exc_class]
else:
raises_args = [kw('Exception', prefix='')]
reason = capture.get('reason')
if reason:
assert len(reason) == 1
reason = KeywordArg(kw('message'), reason[0].clone())
raises_args = [Node(syms.arglist, raises_args + [Comma(), reason])]
raises_args = [LParen()] + raises_args + [RParen()]
capture['fail_stmt'].remove()
try_suite = capture['try_suite'].clone()
with_stmt = Node(
syms.with_stmt,
[
kw('with', prefix=''),
Node(
syms.power,
[
kw('pytest'),
Node(syms.trailer, [Dot(), kw('raises', prefix='')]),
Node(syms.trailer, raises_args),
],
),
Leaf(TOKEN.COLON, ':'),
try_suite,
],
prefix=node.prefix,
)
# Trailing whitespace and any comments after the if statement are captured
# in the prefix for the dedent node. Copy it to the following node.
dedent = capture["dedent"]
next_node = node.next_sibling
# This extra newline avoids syntax errors in some cases (where the try
# statement is at the end of another suite)
# I don't really know why those occur.
# Should clean this stuff up with `black` later.
node.replace([with_stmt, Newline()])
next_node.prefix = dedent.prefix
