def make_dict_comprehension(node, capture, arguments):
"""
dict([(k, v) for k, v in x])
--> {k: v for k, v in x}
PYTHON 2 NOTE:
Where list comprehensions in python 2 set local-scope variables,
dict comprehensions do not!
So this may change the behaviour of your code in subtle ways.
e.g.
>>> a = 5
>>> b = dict([(a, a) for a in (1, 2, 3)])
>>> print(a)
3
>>> a = 5
>>> b = {a: a for a in (1, 2, 3)}
>>> print(a)
5
"""
kv = capture['kv']
key = capture['k']
value = capture['v']
forloop = capture['forloop'][0]
ifpart = capture.get('ifpart') or None
forloop.type = syms.comp_for
if ifpart:
ifpart.type = syms.comp_if
node.replace(
Node(
syms.atom,
[
Leaf(TOKEN.LBRACE, "{"),
Node(
syms.dictsetmaker,
[
key.clone(),
Leaf(TOKEN.COLON, ":"),
value.clone(),
forloop.clone(),
],
prefix=kv.parent.prefix,
),
Leaf(TOKEN.RBRACE, "}", prefix=kv.parent.get_suffix()),
],
prefix=node.prefix,
))
def remove_resolve_only_args(node: LN, capture: Capture,
filename: Filename) -> Optional[LN]:
"""
This one doesn't matter to us since there are no occurences of @resolve_only_args.
XXX: don't use this
"""
# Remove the decorator
decorator = capture.get('decorator')
if decorator is None:
decorators = capture.get('decorators')
else:
decorators = Node(children=[decorator])
if decorator is not None:
for d in decorator.children:
print(f'{d}')
print(f'Child count: {len(decorator.children)}')
if Leaf(1, 'resolve_only_args') in decorator.children:
pos = decorator.children.index(Leaf(1, 'resolve_only_args'))
print(f'at pos: {pos}')
print(f'{decorator.children[pos]}')
decorator.remove()
# Add 'info' to the parameter list
print(capture.get('resolver'))
param = capture.get('param')
if param is not None:
print(param)
return node
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 replace_unicode_methods(node, capture, arguments):
# remove any existing __str__ method
b = find_binding("__str__", capture['suite'])
if b and b.type == syms.funcdef:
b.remove()
# rename __unicode__ to __str__
funcname = capture['funcname'].clone()
funcname.value = '__str__'
capture['funcname'].replace(funcname)
# Add a six import
touch_import(None, "six", node)
# Decorate the class with `@six.python_2_unicode_compatible`
classdef = node.clone()
classdef.prefix = ''
decorated = Node(
syms.decorated,
[
Node(
syms.decorator,
[
Leaf(TOKEN.AT, '@', prefix=node.prefix),
Node(
syms.dotted_name,
[
Name('six'),
Dot(),
Name('python_2_unicode_compatible')
],
),
Newline(),
],
prefix=node.prefix,
),
classdef,
],
prefix=node.prefix,
)
node.replace(decorated)
def simplify_none_operand(node, capture, arguments):
"""
a != None
--> a is not None
"""
op = capture['op'][0]
print(op)
if op.type == TOKEN.EQEQUAL:
op.replace(kw('is'))
else:
op.replace(Node(syms.comp_op, [kw('is'), kw('not')]))
def make_set_comprehension(node, capture, arguments):
arg = capture['arg']
forloop = capture['forloop'][0]
ifpart = capture.get('ifpart') or None
forloop.type = syms.comp_for
if ifpart:
ifpart.type = syms.comp_if
node.replace(
Node(
syms.atom,
[
Leaf(TOKEN.LBRACE, "{"),
Node(
syms.dictsetmaker,
[arg.clone(), forloop.clone()],
prefix=arg.parent.prefix,
),
Leaf(TOKEN.RBRACE, "}", prefix=arg.parent.get_suffix()),
],
prefix=node.prefix,
))
def handle_assertraises(node, capture, arguments):
"""
with self.assertRaises(x):
--> with pytest.raises(x):
self.assertRaises(ValueError, func, arg1)
--> pytest.raises(ValueError, func, arg1)
"""
capture['attr1'].replace(kw('pytest', prefix=capture['attr1'].prefix))
capture['attr2'].replace(
Node(syms.trailer, [Dot(), kw('raises', prefix='')]))
# Adds a 'import pytest' if there wasn't one already
touch_import(None, "pytest", node)
def Assert(test, message=None, **kwargs):
"""Build an assertion statement"""
if not isinstance(test, list):
test = [test]
test[0].prefix = " "
if message is not None:
if not isinstance(message, list):
message = [message]
message.insert(0, Comma())
message[1].prefix = " "
return Node(
syms.assert_stmt,
[Leaf(TOKEN.NAME, "assert")] + test + (message or []),
**kwargs,
)
def handle_assert_regex(node, capture, arguments):
"""
self.assertRegex(text, pattern, msg)
--> assert re.search(pattern, text), msg
self.assertNotRegex(text, pattern, msg)
--> assert not re.search(pattern, text), msg
"""
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()
arguments[0].prefix = " "
arguments[1].prefix = ""
# Adds a 'import re' if there wasn't one already
touch_import(None, "re", node)
op_tokens = []
if invert:
op_tokens.append(keyword("not"))
assert_test_nodes = [
Node(
syms.power,
op_tokens + Attr(keyword("re"), keyword("search", prefix="")) +
[ArgList([arguments[1], Comma(), arguments[0]])],
)
]
return Assert(assert_test_nodes,
message.clone() if message else None,
prefix=node.prefix)
def handle_assertraises(node, capture, arguments):
"""
with self.assertRaises(x):
--> with pytest.raises(x):
self.assertRaises(ValueError, func, arg1)
--> pytest.raises(ValueError, func, arg1)
"""
capture["self_attr"].replace(
keyword("pytest", prefix=capture["self_attr"].prefix))
capture["raises_attr"].replace(
Node(syms.trailer, [Dot(), keyword("raises", prefix="")]))
# Let's remove the msg= keyword argument if found
for child in node.children:
if child.type != syms.trailer:
continue
for tchild in child.children:
if tchild.type != syms.arglist:
continue
previous_argument = None
for argument in tchild.children:
if isinstance(argument, Leaf):
previous_argument = argument
continue
if isinstance(argument, Node):
if argument.type != syms.argument:
previous_argument = argument
continue
for leaf in argument.leaves():
if leaf.value == "msg":
argument.remove()
if previous_argument.value == ",":
# previous_argument is a comma, remove it.
previous_argument.remove()
# Adds a 'import pytest' if there wasn't one already
touch_import(None, "pytest", node)
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 kw(name, **kwargs):
"""
A helper to produce keyword nodes
"""
kwargs.setdefault('prefix', ' ')
return Leaf(TOKEN.NAME, name, **kwargs)
OPERATOR_INVERSIONS = {
'==': Leaf(TOKEN.NOTEQUAL, '!=', prefix=' '),
'!=': Leaf(TOKEN.EQEQUAL, '==', prefix=' '),
'<': Leaf(TOKEN.GREATEREQUAL, '>=', prefix=' '),
'>': Leaf(TOKEN.LESSEQUAL, '<=', prefix=' '),
'<=': Leaf(TOKEN.GREATER, '>', prefix=' '),
'>=': Leaf(TOKEN.LESS, '<', prefix=' '),
'in': Node(syms.comp_op, [kw('not'), kw('in')], prefix=' '),
'not in': kw('in'),
'is': Node(syms.comp_op, [kw('is'), kw('not')], prefix=' '),
'is not': kw('is'),
}
def invert_operator(op):
return OPERATOR_INVERSIONS[str(op).strip()].clone()
def simplify_not_operators(node, capture, arguments):
"""
not a == b
--> a != b
def handle_assertraises_regex(node, capture, arguments):
"""
with self.assertRaisesRegex(x, "regex match"):
--> with pytest.raises(x, match="regex match"):
self.assertRaises(ValueError, "regex match", func, arg1)
--> pytest.raises(ValueError, func, arg1, match="regex match")
"""
capture["self_attr"].replace(
keyword("pytest", prefix=capture["self_attr"].prefix))
capture["raises_attr"].replace(
Node(syms.trailer, [Dot(), keyword("raises", prefix="")]))
# Let's remove the msg= keyword argument if found and replace the second argument
# with a match keyword argument
regex_match = None
for child in node.children:
if child.type != syms.trailer:
continue
for tchild in child.children:
if tchild.type != syms.arglist:
continue
previous_argument = None
argnum = 0
for argument in list(tchild.children):
if isinstance(argument, Leaf):
if argument.value != ",":
argnum += 1
else:
previous_argument = argument
continue
if argnum != 2:
previous_argument = argument
continue
if argnum == 2:
regex_match = Node(
syms.argument,
[
Leaf(TOKEN.NAME, "match"),
Leaf(TOKEN.EQUAL, "="),
Leaf(TOKEN.STRING, argument.value),
],
prefix=" ",
)
argument.remove()
if previous_argument and previous_argument.value == ",":
previous_argument.remove()
previous_argument = None
continue
if isinstance(argument, Node):
if argument.type != syms.argument:
previous_argument = argument
continue
for leaf in argument.leaves():
if leaf.value == "msg":
argument.remove()
if previous_argument and previous_argument.value == ",":
# previous_argument is a comma, remove it.
previous_argument.remove()
if regex_match:
regex_match_added = False
for child in node.children:
if regex_match_added:
break
if child.type != syms.trailer:
continue
for tchild in child.children:
if tchild.type != syms.arglist:
continue
tchild.children.append(Leaf(TOKEN.COMMA, ","))
tchild.children.append(regex_match)
regex_match_added = True
break
# Adds a 'import pytest' if there wasn't one already
touch_import(None, "pytest", 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
def invert_condition(condition, nested=False):
"""
Inverts a boolean expression, e.g.:
a == b
--> a != b
a > b
--> a <= b
a or b
--> not (a or b)
(a or b)
--> not (a or b)
(a and b)
--> not (a and b)
(a == b and c != d)
--> (a != b or c == d)
a if b else c
--> not (a if b else c)
"""
if condition.type == syms.comparison:
a, op, b = condition.children
op = condition.children[1]
if op.type == syms.comp_op:
if (op.children[0].value, op.children[1].value) == ("is", "not"):
return Node(syms.comparison, [a.clone(), kw("is"), b.clone()])
elif (op.children[0].value, op.children[1].value) == ("not", "in"):
return Node(syms.comparison, [a.clone(), kw("in"), b.clone()])
else:
raise NotImplementedError(f"unknown comp_op: {op!r}")
else:
inversions = {
"is": Node(syms.comp_op, [kw("is"), kw("not")], prefix=" "),
"in": Node(syms.comp_op, [kw("not"), kw("in")], prefix=" "),
"==": Leaf(TOKEN.NOTEQUAL, "!=", prefix=" "),
"!=": Leaf(TOKEN.EQEQUAL, "==", prefix=" "),
">": Leaf(TOKEN.LESSEQUAL, "<=", prefix=" "),
"<": Leaf(TOKEN.GREATEREQUAL, ">=", prefix=" "),
"<=": Leaf(TOKEN.GREATER, ">", prefix=" "),
">=": Leaf(TOKEN.LESS, "<", prefix=" "),
}
return Node(syms.comparison,
[a.clone(), inversions[op.value],
b.clone()])
elif condition.type == syms.not_test:
# `not x` --> just remove the `not`
return condition.children[1].clone()
elif condition.type in (syms.and_test, syms.or_test):
# Tricky one.
# (a != b and c != d)
# which should this become?
# --> (a == b or c == d)
# --> not (a != b and c != d)
# Seems somewhat context dependent. Basically we compute both, and then
# decide based on which has the least negations.
simply_inverted = Node(
syms.not_test,
[kw("not"),
parenthesize_if_not_already(condition.clone())])
if condition.type == syms.and_test:
children = [invert_condition(condition.children[0], nested=True)]
for child in condition.children[2::2]:
children.extend(
[kw('or'), invert_condition(child, nested=True)])
complex_inverted = Node(syms.or_test, children)
if nested:
# We're inside an outer 'and' test, and 'or' has lower precedence.
# so we need to parenthesize to ensure the expression is correct
complex_inverted = _parenthesize(complex_inverted)
else:
children = [invert_condition(condition.children[0], nested=True)]
for child in condition.children[2::2]:
children.extend(
[kw('and'),
invert_condition(child, nested=True)])
complex_inverted = Node(syms.and_test, children)
return min([simply_inverted, complex_inverted], key=_num_negations)
if len(str(simply_inverted)) < len(str(complex_inverted)):
return simply_inverted
else:
return complex_inverted
else:
return Node(
syms.not_test,
[kw("not"),
parenthesize_if_not_already(condition.clone())])