def visit_UnaryOperation(self, node: cst.UnaryOperation) -> None:
if m.matches(node, m.UnaryOperation(operator=m.Not(), expression=m.Name())):
# Eg: "not x".
expression: cst.Name = cast(cst.Name, node.expression)
if self._is_optional_type(expression):
replacement_comparison = self._gen_comparison_to_none(
variable_name=expression.value, operator=cst.Is()
)
self.report(node, replacement=replacement_comparison)
def convert_none_cmp(
self, _, updated_node: cst.ComparisonTarget
) -> Union[cst.ComparisonTarget, cst.RemovalSentinel]:
original_op = cst.ensure_type(updated_node.operator, cst.Equal)
return updated_node.with_changes(operator=cst.Is(
whitespace_after=original_op.whitespace_after,
whitespace_before=original_op.whitespace_before,
))
def alter_operator(
self, original_op: Union[cst.Equal,
cst.NotEqual]) -> Union[cst.Is, cst.IsNot]:
return (cst.IsNot(
whitespace_before=original_op.whitespace_before,
whitespace_after=original_op.whitespace_after,
) if isinstance(original_op, cst.NotEqual) else cst.Is(
whitespace_before=original_op.whitespace_before,
whitespace_after=original_op.whitespace_after,
))
class ComparisonTest(CSTNodeTest):
@data_provider((
# Simple comparison statements
(
cst.Comparison(
cst.Name("foo"),
(cst.ComparisonTarget(cst.LessThan(), cst.Integer("5")), ),
),
"foo < 5",
),
(
cst.Comparison(
cst.Name("foo"),
(cst.ComparisonTarget(cst.NotEqual(), cst.Integer("5")), ),
),
"foo != 5",
),
(
cst.Comparison(
cst.Name("foo"),
(cst.ComparisonTarget(cst.Is(), cst.Name("True")), )),
"foo is True",
),
(
cst.Comparison(
cst.Name("foo"),
(cst.ComparisonTarget(cst.IsNot(), cst.Name("False")), ),
),
"foo is not False",
),
(
cst.Comparison(
cst.Name("foo"),
(cst.ComparisonTarget(cst.In(), cst.Name("bar")), )),
"foo in bar",
),
(
cst.Comparison(
cst.Name("foo"),
(cst.ComparisonTarget(cst.NotIn(), cst.Name("bar")), ),
),
"foo not in bar",
),
# Comparison with parens
(
cst.Comparison(
lpar=(cst.LeftParen(), ),
left=cst.Name("foo"),
comparisons=(cst.ComparisonTarget(
operator=cst.NotIn(), comparator=cst.Name("bar")), ),
rpar=(cst.RightParen(), ),
),
"(foo not in bar)",
),
(
cst.Comparison(
left=cst.Name("a",
lpar=(cst.LeftParen(), ),
rpar=(cst.RightParen(), )),
comparisons=(
cst.ComparisonTarget(
operator=cst.Is(
whitespace_before=cst.SimpleWhitespace(""),
whitespace_after=cst.SimpleWhitespace(""),
),
comparator=cst.Name("b",
lpar=(cst.LeftParen(), ),
rpar=(cst.RightParen(), )),
),
cst.ComparisonTarget(
operator=cst.Is(
whitespace_before=cst.SimpleWhitespace(""),
whitespace_after=cst.SimpleWhitespace(""),
),
comparator=cst.Name("c",
lpar=(cst.LeftParen(), ),
rpar=(cst.RightParen(), )),
),
),
),
"(a)is(b)is(c)",
),
# Valid expressions that look like they shouldn't parse
(
cst.Comparison(
left=cst.Integer("5"),
comparisons=(cst.ComparisonTarget(
operator=cst.NotIn(
whitespace_before=cst.SimpleWhitespace("")),
comparator=cst.Name("bar"),
), ),
),
"5not in bar",
),
# Validate that spacing works properly
(
cst.Comparison(
lpar=(cst.LeftParen(
whitespace_after=cst.SimpleWhitespace(" ")), ),
left=cst.Name("foo"),
comparisons=(cst.ComparisonTarget(
operator=cst.NotIn(
whitespace_before=cst.SimpleWhitespace(" "),
whitespace_between=cst.SimpleWhitespace(" "),
whitespace_after=cst.SimpleWhitespace(" "),
),
comparator=cst.Name("bar"),
), ),
rpar=(cst.RightParen(
whitespace_before=cst.SimpleWhitespace(" ")), ),
),
"( foo not in bar )",
),
# Do some complex nodes
(
cst.Comparison(
left=cst.Name("baz"),
comparisons=(cst.ComparisonTarget(
operator=cst.Equal(),
comparator=cst.Comparison(
lpar=(cst.LeftParen(), ),
left=cst.Name("foo"),
comparisons=(cst.ComparisonTarget(
operator=cst.NotIn(),
comparator=cst.Name("bar")), ),
rpar=(cst.RightParen(), ),
),
), ),
),
"baz == (foo not in bar)",
CodeRange((1, 0), (1, 23)),
),
(
cst.Comparison(
left=cst.Name("a"),
comparisons=(
cst.ComparisonTarget(operator=cst.GreaterThan(),
comparator=cst.Name("b")),
cst.ComparisonTarget(operator=cst.GreaterThan(),
comparator=cst.Name("c")),
),
),
"a > b > c",
CodeRange((1, 0), (1, 9)),
),
# Is safe to use with word operators if it's leading/trailing children are
(
cst.IfExp(
body=cst.Comparison(
left=cst.Name("a"),
comparisons=(cst.ComparisonTarget(
operator=cst.GreaterThan(),
comparator=cst.Name(
"b",
lpar=(cst.LeftParen(), ),
rpar=(cst.RightParen(), ),
),
), ),
),
test=cst.Comparison(
left=cst.Name("c",
lpar=(cst.LeftParen(), ),
rpar=(cst.RightParen(), )),
comparisons=(cst.ComparisonTarget(
operator=cst.GreaterThan(),
comparator=cst.Name("d")), ),
),
orelse=cst.Name("e"),
whitespace_before_if=cst.SimpleWhitespace(""),
whitespace_after_if=cst.SimpleWhitespace(""),
),
"a > (b)if(c) > d else e",
),
# is safe to use with word operators if entirely surrounded in parenthesis
(
cst.IfExp(
body=cst.Name("a"),
test=cst.Comparison(
left=cst.Name("b"),
comparisons=(cst.ComparisonTarget(
operator=cst.GreaterThan(),
comparator=cst.Name("c")), ),
lpar=(cst.LeftParen(), ),
rpar=(cst.RightParen(), ),
),
orelse=cst.Name("d"),
whitespace_after_if=cst.SimpleWhitespace(""),
whitespace_before_else=cst.SimpleWhitespace(""),
),
"a if(b > c)else d",
),
))
def test_valid(self,
node: cst.CSTNode,
code: str,
position: Optional[CodeRange] = None) -> None:
self.validate_node(node,
code,
parse_expression,
expected_position=position)
@data_provider((
(
lambda: cst.Comparison(
cst.Name("foo"),
(cst.ComparisonTarget(cst.LessThan(), cst.Integer("5")), ),
lpar=(cst.LeftParen(), ),
),
"left paren without right paren",
),
(
lambda: cst.Comparison(
cst.Name("foo"),
(cst.ComparisonTarget(cst.LessThan(), cst.Integer("5")), ),
rpar=(cst.RightParen(), ),
),
"right paren without left paren",
),
(
lambda: cst.Comparison(cst.Name("foo"), ()),
"at least one ComparisonTarget",
),
(
lambda: cst.Comparison(
left=cst.Name("foo"),
comparisons=(cst.ComparisonTarget(
operator=cst.NotIn(whitespace_before=cst.SimpleWhitespace(
"")),
comparator=cst.Name("bar"),
), ),
),
"at least one space around comparison operator",
),
(
lambda: cst.Comparison(
left=cst.Name("foo"),
comparisons=(cst.ComparisonTarget(
operator=cst.NotIn(whitespace_after=cst.SimpleWhitespace(
"")),
comparator=cst.Name("bar"),
), ),
),
"at least one space around comparison operator",
),
# multi-target comparisons
(
lambda: cst.Comparison(
left=cst.Name("a"),
comparisons=(
cst.ComparisonTarget(operator=cst.Is(),
comparator=cst.Name("b")),
cst.ComparisonTarget(
operator=cst.Is(whitespace_before=cst.SimpleWhitespace(
"")),
comparator=cst.Name("c"),
),
),
),
"at least one space around comparison operator",
),
(
lambda: cst.Comparison(
left=cst.Name("a"),
comparisons=(
cst.ComparisonTarget(operator=cst.Is(),
comparator=cst.Name("b")),
cst.ComparisonTarget(
operator=cst.Is(whitespace_after=cst.SimpleWhitespace(
"")),
comparator=cst.Name("c"),
),
),
),
"at least one space around comparison operator",
),
# whitespace around the comparision itself
# a ifb > c else d
(
lambda: cst.IfExp(
body=cst.Name("a"),
test=cst.Comparison(
left=cst.Name("b"),
comparisons=(cst.
ComparisonTarget(operator=cst.GreaterThan(),
comparator=cst.Name("c")), ),
),
orelse=cst.Name("d"),
whitespace_after_if=cst.SimpleWhitespace(""),
),
"Must have at least one space after 'if' keyword.",
),
# a if b > celse d
(
lambda: cst.IfExp(
body=cst.Name("a"),
test=cst.Comparison(
left=cst.Name("b"),
comparisons=(cst.
ComparisonTarget(operator=cst.GreaterThan(),
comparator=cst.Name("c")), ),
),
orelse=cst.Name("d"),
whitespace_before_else=cst.SimpleWhitespace(""),
),
"Must have at least one space before 'else' keyword.",
),
))
def test_invalid(self, get_node: Callable[[], cst.CSTNode],
expected_re: str) -> None:
self.assert_invalid(get_node, expected_re)
class SimpleCompTest(CSTNodeTest):
@data_provider(
[
# simple GeneratorExp
{
"node": cst.GeneratorExp(
cst.Name("a"), cst.CompFor(target=cst.Name("b"), iter=cst.Name("c"))
),
"code": "(a for b in c)",
"parser": parse_expression,
"expected_position": CodeRange.create((1, 1), (1, 13)),
},
# simple ListComp
{
"node": cst.ListComp(
cst.Name("a"), cst.CompFor(target=cst.Name("b"), iter=cst.Name("c"))
),
"code": "[a for b in c]",
"parser": parse_expression,
"expected_position": CodeRange.create((1, 0), (1, 14)),
},
# simple SetComp
{
"node": cst.SetComp(
cst.Name("a"), cst.CompFor(target=cst.Name("b"), iter=cst.Name("c"))
),
"code": "{a for b in c}",
"parser": parse_expression,
},
# async GeneratorExp
{
"node": cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
asynchronous=cst.Asynchronous(),
),
),
"code": "(a async for b in c)",
"parser": parse_expression,
},
# a generator doesn't have to own it's own parenthesis
{
"node": cst.Call(
cst.Name("func"),
[
cst.Arg(
cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
lpar=[],
rpar=[],
)
)
],
),
"code": "func(a for b in c)",
"parser": parse_expression,
},
# add a few 'if' clauses
{
"node": cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
ifs=[
cst.CompIf(cst.Name("d")),
cst.CompIf(cst.Name("e")),
cst.CompIf(cst.Name("f")),
],
),
),
"code": "(a for b in c if d if e if f)",
"parser": parse_expression,
"expected_position": CodeRange.create((1, 1), (1, 28)),
},
# nested/inner for-in clause
{
"node": cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
inner_for_in=cst.CompFor(
target=cst.Name("d"), iter=cst.Name("e")
),
),
),
"code": "(a for b in c for d in e)",
"parser": parse_expression,
},
# nested/inner for-in clause with an 'if' clause
{
"node": cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
ifs=[cst.CompIf(cst.Name("d"))],
inner_for_in=cst.CompFor(
target=cst.Name("e"), iter=cst.Name("f")
),
),
),
"code": "(a for b in c if d for e in f)",
"parser": parse_expression,
},
# custom whitespace
{
"node": cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
ifs=[
cst.CompIf(
cst.Name("d"),
whitespace_before=cst.SimpleWhitespace("\t"),
whitespace_before_test=cst.SimpleWhitespace("\t\t"),
)
],
whitespace_before=cst.SimpleWhitespace(" "),
whitespace_after_for=cst.SimpleWhitespace(" "),
whitespace_before_in=cst.SimpleWhitespace(" "),
whitespace_after_in=cst.SimpleWhitespace(" "),
),
lpar=[cst.LeftParen(whitespace_after=cst.SimpleWhitespace("\f"))],
rpar=[
cst.RightParen(whitespace_before=cst.SimpleWhitespace("\f\f"))
],
),
"code": "(\fa for b in c\tif\t\td\f\f)",
"parser": parse_expression,
"expected_position": CodeRange.create((1, 2), (1, 30)),
},
# custom whitespace around ListComp's brackets
{
"node": cst.ListComp(
cst.Name("a"),
cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
lbracket=cst.LeftSquareBracket(
whitespace_after=cst.SimpleWhitespace("\t")
),
rbracket=cst.RightSquareBracket(
whitespace_before=cst.SimpleWhitespace("\t\t")
),
lpar=[cst.LeftParen(whitespace_after=cst.SimpleWhitespace("\f"))],
rpar=[
cst.RightParen(whitespace_before=cst.SimpleWhitespace("\f\f"))
],
),
"code": "(\f[\ta for b in c\t\t]\f\f)",
"parser": parse_expression,
"expected_position": CodeRange.create((1, 2), (1, 19)),
},
# custom whitespace around SetComp's braces
{
"node": cst.SetComp(
cst.Name("a"),
cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
lbrace=cst.LeftCurlyBrace(
whitespace_after=cst.SimpleWhitespace("\t")
),
rbrace=cst.RightCurlyBrace(
whitespace_before=cst.SimpleWhitespace("\t\t")
),
lpar=[cst.LeftParen(whitespace_after=cst.SimpleWhitespace("\f"))],
rpar=[
cst.RightParen(whitespace_before=cst.SimpleWhitespace("\f\f"))
],
),
"code": "(\f{\ta for b in c\t\t}\f\f)",
"parser": parse_expression,
},
# no whitespace between elements
{
"node": cst.GeneratorExp(
cst.Name("a", lpar=[cst.LeftParen()], rpar=[cst.RightParen()]),
cst.CompFor(
target=cst.Name(
"b", lpar=[cst.LeftParen()], rpar=[cst.RightParen()]
),
iter=cst.Name(
"c", lpar=[cst.LeftParen()], rpar=[cst.RightParen()]
),
ifs=[
cst.CompIf(
cst.Name(
"d", lpar=[cst.LeftParen()], rpar=[cst.RightParen()]
),
whitespace_before=cst.SimpleWhitespace(""),
whitespace_before_test=cst.SimpleWhitespace(""),
)
],
inner_for_in=cst.CompFor(
target=cst.Name(
"e", lpar=[cst.LeftParen()], rpar=[cst.RightParen()]
),
iter=cst.Name(
"f", lpar=[cst.LeftParen()], rpar=[cst.RightParen()]
),
whitespace_before=cst.SimpleWhitespace(""),
whitespace_after_for=cst.SimpleWhitespace(""),
whitespace_before_in=cst.SimpleWhitespace(""),
whitespace_after_in=cst.SimpleWhitespace(""),
),
whitespace_before=cst.SimpleWhitespace(""),
whitespace_after_for=cst.SimpleWhitespace(""),
whitespace_before_in=cst.SimpleWhitespace(""),
whitespace_after_in=cst.SimpleWhitespace(""),
),
lpar=[cst.LeftParen()],
rpar=[cst.RightParen()],
),
"code": "((a)for(b)in(c)if(d)for(e)in(f))",
"parser": parse_expression,
"expected_position": CodeRange.create((1, 1), (1, 31)),
},
# no whitespace before/after GeneratorExp is valid
{
"node": cst.Comparison(
cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
),
[
cst.ComparisonTarget(
cst.Is(
whitespace_before=cst.SimpleWhitespace(""),
whitespace_after=cst.SimpleWhitespace(""),
),
cst.GeneratorExp(
cst.Name("d"),
cst.CompFor(target=cst.Name("e"), iter=cst.Name("f")),
),
)
],
),
"code": "(a for b in c)is(d for e in f)",
"parser": parse_expression,
},
# no whitespace before/after ListComp is valid
{
"node": cst.Comparison(
cst.ListComp(
cst.Name("a"),
cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
),
[
cst.ComparisonTarget(
cst.Is(
whitespace_before=cst.SimpleWhitespace(""),
whitespace_after=cst.SimpleWhitespace(""),
),
cst.ListComp(
cst.Name("d"),
cst.CompFor(target=cst.Name("e"), iter=cst.Name("f")),
),
)
],
),
"code": "[a for b in c]is[d for e in f]",
"parser": parse_expression,
},
# no whitespace before/after SetComp is valid
{
"node": cst.Comparison(
cst.SetComp(
cst.Name("a"),
cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
),
[
cst.ComparisonTarget(
cst.Is(
whitespace_before=cst.SimpleWhitespace(""),
whitespace_after=cst.SimpleWhitespace(""),
),
cst.SetComp(
cst.Name("d"),
cst.CompFor(target=cst.Name("e"), iter=cst.Name("f")),
),
)
],
),
"code": "{a for b in c}is{d for e in f}",
"parser": parse_expression,
},
]
)
def test_valid(self, **kwargs: Any) -> None:
self.validate_node(**kwargs)
@data_provider(
(
(
lambda: cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
lpar=[cst.LeftParen(), cst.LeftParen()],
rpar=[cst.RightParen()],
),
"unbalanced parens",
),
(
lambda: cst.ListComp(
cst.Name("a"),
cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
lpar=[cst.LeftParen(), cst.LeftParen()],
rpar=[cst.RightParen()],
),
"unbalanced parens",
),
(
lambda: cst.SetComp(
cst.Name("a"),
cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
lpar=[cst.LeftParen(), cst.LeftParen()],
rpar=[cst.RightParen()],
),
"unbalanced parens",
),
(
lambda: cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
whitespace_before=cst.SimpleWhitespace(""),
),
),
"Must have at least one space before 'for' keyword.",
),
(
lambda: cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
asynchronous=cst.Asynchronous(),
whitespace_before=cst.SimpleWhitespace(""),
),
),
"Must have at least one space before 'async' keyword.",
),
(
lambda: cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
whitespace_after_for=cst.SimpleWhitespace(""),
),
),
"Must have at least one space after 'for' keyword.",
),
(
lambda: cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
whitespace_before_in=cst.SimpleWhitespace(""),
),
),
"Must have at least one space before 'in' keyword.",
),
(
lambda: cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
whitespace_after_in=cst.SimpleWhitespace(""),
),
),
"Must have at least one space after 'in' keyword.",
),
(
lambda: cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
ifs=[
cst.CompIf(
cst.Name("d"),
whitespace_before=cst.SimpleWhitespace(""),
)
],
),
),
"Must have at least one space before 'if' keyword.",
),
(
lambda: cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
ifs=[
cst.CompIf(
cst.Name("d"),
whitespace_before_test=cst.SimpleWhitespace(""),
)
],
),
),
"Must have at least one space after 'if' keyword.",
),
(
lambda: cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
inner_for_in=cst.CompFor(
target=cst.Name("d"),
iter=cst.Name("e"),
whitespace_before=cst.SimpleWhitespace(""),
),
),
),
"Must have at least one space before 'for' keyword.",
),
(
lambda: cst.GeneratorExp(
cst.Name("a"),
cst.CompFor(
target=cst.Name("b"),
iter=cst.Name("c"),
inner_for_in=cst.CompFor(
target=cst.Name("d"),
iter=cst.Name("e"),
asynchronous=cst.Asynchronous(),
whitespace_before=cst.SimpleWhitespace(""),
),
),
),
"Must have at least one space before 'async' keyword.",
),
)
)
def test_invalid(
self, get_node: Callable[[], cst.CSTNode], expected_re: str
) -> None:
self.assert_invalid(get_node, expected_re)
def convert_none_cmp(self, _, updated_node):
"""Inspired by Pybetter."""
return updated_node.with_changes(operator=cst.Is())