Allow no space around an ifexp in certain cases

For example in `_ if _ else""if _ else _`.

Raised in #930. Also fixes #854.
This commit is contained in:
Zsolt Dollenstein 2023-05-28 10:29:29 +01:00
parent f45ba63ca6
commit aa5167189c
No known key found for this signature in database
GPG key ID: 09F75E762C27AAD9
2 changed files with 41 additions and 0 deletions

View file

@ -2512,6 +2512,12 @@ class IfExp(BaseExpression):
#: Whitespace after the ``else`` keyword, but before the ``orelse`` expression.
whitespace_after_else: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
if position == ExpressionPosition.RIGHT:
return self.body._safe_to_use_with_word_operator(position)
else:
return self.orelse._safe_to_use_with_word_operator(position)
def _validate(self) -> None:
# Paren validation and such
super(IfExp, self)._validate()

View file

@ -52,6 +52,41 @@ class IfExpTest(CSTNodeTest):
"(foo)if(bar)else(baz)",
CodeRange((1, 0), (1, 21)),
),
(
cst.IfExp(
body=cst.Name("foo"),
whitespace_before_if=cst.SimpleWhitespace(" "),
whitespace_after_if=cst.SimpleWhitespace(" "),
test=cst.Name("bar"),
whitespace_before_else=cst.SimpleWhitespace(" "),
whitespace_after_else=cst.SimpleWhitespace(""),
orelse=cst.IfExp(
body=cst.SimpleString("''"),
whitespace_before_if=cst.SimpleWhitespace(""),
test=cst.Name("bar"),
orelse=cst.Name("baz"),
),
),
"foo if bar else''if bar else baz",
CodeRange((1, 0), (1, 32)),
),
(
cst.GeneratorExp(
elt=cst.IfExp(
body=cst.Name("foo"),
test=cst.Name("bar"),
orelse=cst.SimpleString("''"),
whitespace_after_else=cst.SimpleWhitespace(""),
),
for_in=cst.CompFor(
target=cst.Name("_"),
iter=cst.Name("_"),
whitespace_before=cst.SimpleWhitespace(""),
),
),
"(foo if bar else''for _ in _)",
CodeRange((1, 1), (1, 28)),
),
# Make sure that spacing works
(
cst.IfExp(