Merge pull request #249 from ncaklovic/master
Support for conditional expression (if-expression)
This commit is contained in:
31
ASTNode.h
31
ASTNode.h
@@ -16,7 +16,7 @@ public:
|
||||
NODE_CONVERT, NODE_KEYWORD, NODE_RAISE, NODE_EXEC, NODE_BLOCK,
|
||||
NODE_COMPREHENSION, NODE_LOADBUILDCLASS, NODE_AWAITABLE,
|
||||
NODE_FORMATTEDVALUE, NODE_JOINEDSTR, NODE_CONST_MAP,
|
||||
NODE_ANNOTATED_VAR, NODE_CHAINSTORE,
|
||||
NODE_ANNOTATED_VAR, NODE_CHAINSTORE, NODE_TERNARY,
|
||||
|
||||
// Empty node types
|
||||
NODE_LOCALS,
|
||||
@@ -699,4 +699,33 @@ private:
|
||||
PycRef<ASTNode> m_type;
|
||||
};
|
||||
|
||||
class ASTTernary : public ASTNode
|
||||
{
|
||||
public:
|
||||
ASTTernary(PycRef<ASTNode> if_block, PycRef<ASTNode> if_expr, PycRef<ASTNode> else_expr)
|
||||
: ASTNode(NODE_TERNARY),
|
||||
m_if_block(std::move(if_block)),
|
||||
m_if_expr(std::move(if_expr)),
|
||||
m_else_expr(std::move(else_expr))
|
||||
{
|
||||
}
|
||||
const PycRef<ASTNode>& if_block() const noexcept
|
||||
{
|
||||
return m_if_block;
|
||||
}
|
||||
const PycRef<ASTNode>& if_expr() const noexcept
|
||||
{
|
||||
return m_if_expr;
|
||||
}
|
||||
const PycRef<ASTNode>& else_expr() const noexcept
|
||||
{
|
||||
return m_else_expr;
|
||||
}
|
||||
|
||||
private:
|
||||
PycRef<ASTNode> m_if_block; // contains "condition" and "negative"
|
||||
PycRef<ASTNode> m_if_expr;
|
||||
PycRef<ASTNode> m_else_expr;
|
||||
};
|
||||
|
||||
#endif
|
||||
|
76
ASTree.cpp
76
ASTree.cpp
@@ -29,6 +29,43 @@ static bool printDocstringAndGlobals = false;
|
||||
/* Use this to keep track of whether we need to print a class or module docstring */
|
||||
static bool printClassDocstring = true;
|
||||
|
||||
// shortcut for all top/pop calls
|
||||
static PycRef<ASTNode> StackPopTop(FastStack& stack)
|
||||
{
|
||||
const auto node{ stack.top() };
|
||||
stack.pop();
|
||||
return node;
|
||||
}
|
||||
|
||||
/* compiler generates very, VERY similar byte code for if/else statement block and if-expression
|
||||
* statement
|
||||
* if a: b = 1
|
||||
* else: b = 2
|
||||
* expression:
|
||||
* b = 1 if a else 2
|
||||
* (see for instance https://stackoverflow.com/a/52202007)
|
||||
* here, try to guess if just finished else statement is part of if-expression (ternary operator)
|
||||
* if it is, remove statements from the block and put a ternary node on top of stack
|
||||
*/
|
||||
static void CheckIfExpr(FastStack& stack, PycRef<ASTBlock> curblock)
|
||||
{
|
||||
if (stack.empty())
|
||||
return;
|
||||
if (curblock->nodes().size() < 2)
|
||||
return;
|
||||
auto rit{ curblock->nodes().crbegin() };
|
||||
++rit; // the last is "else" block, the one before should be "if" (could be "for", ...)
|
||||
if ((*rit)->type() != ASTNode::NODE_BLOCK ||
|
||||
(*rit).cast<ASTBlock>()->blktype() != ASTBlock::BLK_IF)
|
||||
return;
|
||||
auto else_expr{ StackPopTop(stack) };
|
||||
curblock->removeLast();
|
||||
auto if_block{ curblock->nodes().back() };
|
||||
auto if_expr{ StackPopTop(stack) };
|
||||
curblock->removeLast();
|
||||
stack.push(new ASTTernary(std::move(if_block), std::move(if_expr), std::move(else_expr)));
|
||||
}
|
||||
|
||||
PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
|
||||
{
|
||||
PycBuffer source(code->code()->value(), code->code()->length());
|
||||
@@ -109,6 +146,8 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
|
||||
curblock->append(prev.cast<ASTNode>());
|
||||
|
||||
prev = curblock;
|
||||
|
||||
CheckIfExpr(stack, curblock);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1363,19 +1402,9 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
|
||||
break;
|
||||
}
|
||||
|
||||
if ((curblock->blktype() == ASTBlock::BLK_WHILE
|
||||
&& !curblock->inited())
|
||||
|| (curblock->blktype() == ASTBlock::BLK_IF
|
||||
&& curblock->size() == 0)) {
|
||||
PycRef<PycObject> fakeint = new PycInt(1);
|
||||
PycRef<ASTNode> truthy = new ASTObject(fakeint);
|
||||
|
||||
stack.push(truthy);
|
||||
break;
|
||||
}
|
||||
|
||||
if (!stack_hist.empty()) {
|
||||
stack = stack_hist.top();
|
||||
if (stack.empty()) // if it's part of if-expression, TOS at the moment is the result of "if" part
|
||||
stack = stack_hist.top();
|
||||
stack_hist.pop();
|
||||
}
|
||||
|
||||
@@ -3256,6 +3285,29 @@ void print_src(PycRef<ASTNode> node, PycModule* mod)
|
||||
print_src(type, mod);
|
||||
}
|
||||
break;
|
||||
case ASTNode::NODE_TERNARY:
|
||||
{
|
||||
/* parenthesis might be needed
|
||||
*
|
||||
* when if-expr is part of numerical expression, ternary has the LOWEST precedence
|
||||
* print(a + b if False else c)
|
||||
* output is c, not a+c (a+b is calculated first)
|
||||
*
|
||||
* but, let's not add parenthesis - to keep the source as close to original as possible in most cases
|
||||
*/
|
||||
PycRef<ASTTernary> ternary = node.cast<ASTTernary>();
|
||||
//fputs("(", pyc_output);
|
||||
print_src(ternary->if_expr(), mod);
|
||||
const auto if_block = ternary->if_block().require_cast<ASTCondBlock>();
|
||||
fputs(" if ", pyc_output);
|
||||
if (if_block->negative())
|
||||
fputs("not ", pyc_output);
|
||||
print_src(if_block->cond(), mod);
|
||||
fputs(" else ", pyc_output);
|
||||
print_src(ternary->else_expr(), mod);
|
||||
//fputs(")", pyc_output);
|
||||
}
|
||||
break;
|
||||
default:
|
||||
fprintf(pyc_output, "<NODE:%d>", node->type());
|
||||
fprintf(stderr, "Unsupported Node type: %d\n", node->type());
|
||||
|
@@ -40,6 +40,11 @@ public:
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
bool empty() const
|
||||
{
|
||||
return m_ptr == -1;
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<PycRef<ASTNode>> m_stack;
|
||||
int m_ptr;
|
||||
|
BIN
tests/compiled/conditional_expressions.3.9.pyc
Normal file
BIN
tests/compiled/conditional_expressions.3.9.pyc
Normal file
Binary file not shown.
42
tests/input/conditional_expressions.py
Normal file
42
tests/input/conditional_expressions.py
Normal file
@@ -0,0 +1,42 @@
|
||||
a = 1
|
||||
result = 'even' if a % 2 == 0 else 'odd'
|
||||
print(result)
|
||||
# odd
|
||||
|
||||
a = 2
|
||||
result = 'even' if a % 2 == 0 else 'odd'
|
||||
print(result)
|
||||
# even
|
||||
|
||||
a = 1
|
||||
result = a * 2 if a % 2 == 0 else a * 3
|
||||
print(result)
|
||||
# 3
|
||||
|
||||
a = 2
|
||||
result = a * 2 if a % 2 == 0 else a * 3
|
||||
print(result)
|
||||
# 4
|
||||
|
||||
a = 1
|
||||
print('even') if a % 2 == 0 else print('odd')
|
||||
# odd
|
||||
|
||||
a = 1
|
||||
|
||||
if a % 2 == 0:
|
||||
print('even')
|
||||
else:
|
||||
print('odd')
|
||||
# odd
|
||||
|
||||
a = -2
|
||||
result = 'negative and even' if a < 0 and a % 2 == 0 else 'positive or odd'
|
||||
print(result)
|
||||
# negative and even
|
||||
|
||||
a = -1
|
||||
result = 'negative and even' if a < 0 and a % 2 == 0 else 'positive or odd'
|
||||
print(result)
|
||||
# positive or odd
|
||||
|
29
tests/tokenized/conditional_expressions.txt
Normal file
29
tests/tokenized/conditional_expressions.txt
Normal file
@@ -0,0 +1,29 @@
|
||||
a = 1 <EOL>
|
||||
result = 'even' if a % 2 == 0 else 'odd' <EOL>
|
||||
print ( result ) <EOL>
|
||||
a = 2 <EOL>
|
||||
result = 'even' if a % 2 == 0 else 'odd' <EOL>
|
||||
print ( result ) <EOL>
|
||||
a = 1 <EOL>
|
||||
result = a * 2 if a % 2 == 0 else a * 3 <EOL>
|
||||
print ( result ) <EOL>
|
||||
a = 2 <EOL>
|
||||
result = a * 2 if a % 2 == 0 else a * 3 <EOL>
|
||||
print ( result ) <EOL>
|
||||
a = 1 <EOL>
|
||||
print ( 'even' ) if a % 2 == 0 else print ( 'odd' ) <EOL>
|
||||
a = 1 <EOL>
|
||||
if a % 2 == 0 : <EOL>
|
||||
<INDENT>
|
||||
print ( 'even' ) <EOL>
|
||||
<OUTDENT>
|
||||
else : <EOL>
|
||||
<INDENT>
|
||||
print ( 'odd' ) <EOL>
|
||||
<OUTDENT>
|
||||
a = - 2 <EOL>
|
||||
result = 'negative and even' if a < 0 and a % 2 == 0 else 'positive or odd' <EOL>
|
||||
print ( result ) <EOL>
|
||||
a = - 1 <EOL>
|
||||
result = 'negative and even' if a < 0 and a % 2 == 0 else 'positive or odd' <EOL>
|
||||
print ( result ) <EOL>
|
Reference in New Issue
Block a user