Skip to content

Commit

Permalink
Merge pull request zrax#181 from Aralox/Issue-124-handle-async-for-GE…
Browse files Browse the repository at this point in the history
…T_AITER_GET_ANEXT

Add support for `async for`
  • Loading branch information
zrax authored Oct 22, 2020
2 parents 378c256 + 7022db1 commit 0540b39
Show file tree
Hide file tree
Showing 9 changed files with 323 additions and 21 deletions.
2 changes: 1 addition & 1 deletion ASTNode.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ const char* ASTBlock::type_str() const
{
static const char* s_type_strings[] = {
"", "if", "else", "elif", "try", "CONTAINER", "except",
"finally", "while", "for", "with",
"finally", "while", "for", "with", "async for"
};
return s_type_strings[blktype()];
}
2 changes: 1 addition & 1 deletion ASTNode.h
Original file line number Diff line number Diff line change
Expand Up @@ -482,7 +482,7 @@ class ASTBlock : public ASTNode {
enum BlkType {
BLK_MAIN, BLK_IF, BLK_ELSE, BLK_ELIF, BLK_TRY,
BLK_CONTAINER, BLK_EXCEPT, BLK_FINALLY,
BLK_WHILE, BLK_FOR, BLK_WITH
BLK_WHILE, BLK_FOR, BLK_WITH, BLK_ASYNCFOR
};

ASTBlock(BlkType blktype, int end = 0, int inited = 0)
Expand Down
87 changes: 68 additions & 19 deletions ASTree.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -729,22 +729,51 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
curblock->append(final.cast<ASTNode>());
isFinally = true;
} else if (curblock->blktype() == ASTBlock::BLK_EXCEPT) {
/* Turn it into an else statement. */
blocks.pop();
PycRef<ASTBlock> prev = curblock;
if (curblock->size() != 0) {
blocks.top()->append(curblock.cast<ASTNode>());

bool isUninitAsyncFor = false;
if (blocks.top()->blktype() == ASTBlock::BLK_CONTAINER) {
auto container = blocks.top();
blocks.pop();
auto asyncForBlock = blocks.top();
isUninitAsyncFor = asyncForBlock->blktype() == ASTBlock::BLK_ASYNCFOR && !asyncForBlock->inited();
if (isUninitAsyncFor) {
auto tryBlock = container->nodes().front().cast<ASTBlock>();
if (!tryBlock->nodes().empty() && tryBlock->blktype() == ASTBlock::BLK_TRY) {
auto store = tryBlock->nodes().front().cast<ASTStore>();
if (store) {
asyncForBlock.cast<ASTIterBlock>()->setIndex(store->dest());
}
}
curblock = blocks.top();
stack = stack_hist.top();
stack_hist.pop();
if (!curblock->inited())
fprintf(stderr, "Error when decompiling 'async for'.\n");
} else {
blocks.push(container);
}
}
curblock = blocks.top();

if (curblock->end() != pos || curblock.cast<ASTContainerBlock>()->hasFinally()) {
PycRef<ASTBlock> elseblk = new ASTBlock(ASTBlock::BLK_ELSE, prev->end());
elseblk->init();
blocks.push(elseblk);
if (!isUninitAsyncFor) {
if (curblock->size() != 0) {
blocks.top()->append(curblock.cast<ASTNode>());
}

curblock = blocks.top();
} else {
stack = stack_hist.top();
stack_hist.pop();

/* Turn it into an else statement. */
if (curblock->end() != pos || curblock.cast<ASTContainerBlock>()->hasFinally()) {
PycRef<ASTBlock> elseblk = new ASTBlock(ASTBlock::BLK_ELSE, prev->end());
elseblk->init();
blocks.push(elseblk);
curblock = blocks.top();
}
else {
stack = stack_hist.top();
stack_hist.pop();
}
}
}

Expand Down Expand Up @@ -820,6 +849,26 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
stack.push(NULL); // We can totally hack this >_>
}
break;
case Pyc::GET_AITER:
{
// Logic similar to FOR_ITER_A
PycRef<ASTNode> iter = stack.top(); // Iterable
stack.pop();

PycRef<ASTBlock> top = blocks.top();
if (top->blktype() == ASTBlock::BLK_WHILE) {
blocks.pop();
PycRef<ASTIterBlock> forblk = new ASTIterBlock(ASTBlock::BLK_ASYNCFOR, top->end(), iter);
blocks.push(forblk.cast<ASTBlock>());
curblock = blocks.top();
stack.push(nullptr);
} else {
fprintf(stderr, "Unsupported use of GET_AITER outside of SETUP_LOOP\n");
}
}
break;
case Pyc::GET_ANEXT:
break;
case Pyc::FORMAT_VALUE_A:
{
auto conversion_flag = static_cast<ASTFormattedValue::ConversionFlag>(operand);
Expand Down Expand Up @@ -1448,8 +1497,6 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
break;
}

PycRef<ASTBlock> tmp;

if (curblock->nodes().size() &&
curblock->nodes().back().type() == ASTNode::NODE_KEYWORD) {
curblock->removeLast();
Expand All @@ -1468,8 +1515,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
fprintf(stderr, "Warning: Stack history is empty, something wrong might have happened\n");
}
}

tmp = curblock;
PycRef<ASTBlock> tmp = curblock;
blocks.pop();

if (!blocks.empty())
Expand All @@ -1490,6 +1536,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)

if (curblock->blktype() == ASTBlock::BLK_TRY
&& tmp->blktype() != ASTBlock::BLK_FOR
&& tmp->blktype() != ASTBlock::BLK_ASYNCFOR
&& tmp->blktype() != ASTBlock::BLK_WHILE) {
stack = stack_hist.top();
stack_hist.pop();
Expand Down Expand Up @@ -1526,7 +1573,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
}
}

if (curblock->blktype() == ASTBlock::BLK_FOR
if ((curblock->blktype() == ASTBlock::BLK_FOR || curblock->blktype() == ASTBlock::BLK_ASYNCFOR)
&& curblock->end() == pos) {
blocks.pop();
blocks.top()->append(curblock.cast<ASTNode>());
Expand Down Expand Up @@ -2205,8 +2252,10 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
stack.pop();
// TODO: Support yielding into a non-null destination
PycRef<ASTNode> value = stack.top();
value->setProcessed();
curblock->append(new ASTReturn(value, ASTReturn::YIELD_FROM));
if (value) {
value->setProcessed();
curblock->append(new ASTReturn(value, ASTReturn::YIELD_FROM));
}
}
break;
case Pyc::YIELD_VALUE:
Expand Down Expand Up @@ -2625,7 +2674,7 @@ void print_src(PycRef<ASTNode> node, PycModule* mod)
fputs(" ", pyc_output);

print_src(blk.cast<ASTCondBlock>()->cond(), mod);
} else if (blk->blktype() == ASTBlock::BLK_FOR) {
} else if (blk->blktype() == ASTBlock::BLK_FOR || blk->blktype() == ASTBlock::BLK_ASYNCFOR) {
fputs(" ", pyc_output);
print_src(blk.cast<ASTIterBlock>()->index(), mod);
fputs(" in ", pyc_output);
Expand Down
Binary file added tests/compiled/async_for.3.7.pyc
Binary file not shown.
75 changes: 75 additions & 0 deletions tests/input/async_for.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
async def a(b, c):
async for b in c:
pass

async def a(b, c):
async for b in c:
continue

async def a(b, c):
async for b in c:
break

async def time_for_some_fun():
async for (x, y) in myfunc(c):
print(123)
if (x == 3):
print('something')
break

for i in regular_for:
var1 = var2 + var3
async for x1 in print:
print('test LOAD_GLOBAL')
async for x2 in inner:
for x3 in regular:
async for x4 in inner2:
async for x5 in inner3:
async for x6 in inner4:
print('ridiculous nesting')

async for (q, w, e, r) in qwer:
u = 1 + 2 + 3

async for x4 in inner2:
async for x5 in inner3:
pass

print('outside loop')

async def tryBlocks():
async for b in c:
stuff
await things(x)
try:
STUFF
except MyException:
running = False
BLOCK_AFTER

try:
async for b in c:
stuff
await things(x)
async for b2 in c2:
stuff2
await things2(x)
try:
STUFF2
except MyException2:
running2 = False
AFTER
try:
STUFF
except MyException:
running = False
BLOCK_AFTER

except MyException2:
OUTEREXCEPT

print ('outside function')

# The following will LOAD_METHOD, not GET_AITER or GET_ANEXT.
# test.__anext__(iter)
# test.__aiter__(iter)
20 changes: 20 additions & 0 deletions tests/input/loop_try_except.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
async def myFunc():
async for b in c:
try:
STUFF
except MyException:
running = False

for b in c:
stuff
try:
STUFF
except MyException:
running = False

while a:
stuff
try:
STUFF
except MyException:
running = False
122 changes: 122 additions & 0 deletions tests/tokenized/async_for.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,122 @@
async def a ( b , c ) : <EOL>
<INDENT>
async for b in c : <EOL>
<INDENT>
pass <EOL>
<OUTDENT>
<OUTDENT>
async def a ( b , c ) : <EOL>
<INDENT>
async for b in c : <EOL>
<INDENT>
continue <EOL>
<OUTDENT>
<OUTDENT>
async def a ( b , c ) : <EOL>
<INDENT>
async for b in c : <EOL>
<INDENT>
break <EOL>
<OUTDENT>
<OUTDENT>
async def time_for_some_fun ( ) : <EOL>
<INDENT>
async for ( x , y ) in myfunc ( c ) : <EOL>
<INDENT>
print ( 123 ) <EOL>
if x == 3 : <EOL>
<INDENT>
print ( 'something' ) <EOL>
break <EOL>
<OUTDENT>
for i in regular_for : <EOL>
<INDENT>
var1 = var2 + var3 <EOL>
async for x1 in print : <EOL>
<INDENT>
print ( 'test LOAD_GLOBAL' ) <EOL>
async for x2 in inner : <EOL>
<INDENT>
for x3 in regular : <EOL>
<INDENT>
async for x4 in inner2 : <EOL>
<INDENT>
async for x5 in inner3 : <EOL>
<INDENT>
async for x6 in inner4 : <EOL>
<INDENT>
print ( 'ridiculous nesting' ) <EOL>
<OUTDENT>
<OUTDENT>
<OUTDENT>
<OUTDENT>
<OUTDENT>
<OUTDENT>
<OUTDENT>
<OUTDENT>
async for ( q , w , e , r ) in qwer : <EOL>
<INDENT>
u = 6 <EOL>
<OUTDENT>
async for x4 in inner2 : <EOL>
<INDENT>
async for x5 in inner3 : <EOL>
<INDENT>
pass <EOL>
<OUTDENT>
<OUTDENT>
print ( 'outside loop' ) <EOL>
<OUTDENT>
async def tryBlocks ( ) : <EOL>
<INDENT>
async for b in c : <EOL>
<INDENT>
stuff <EOL>
await things ( x ) <EOL>
try : <EOL>
<INDENT>
STUFF <EOL>
<OUTDENT>
except MyException : <EOL>
<INDENT>
running = False <EOL>
<OUTDENT>
BLOCK_AFTER <EOL>
<OUTDENT>
try : <EOL>
<INDENT>
async for b in c : <EOL>
<INDENT>
stuff <EOL>
await things ( x ) <EOL>
async for b2 in c2 : <EOL>
<INDENT>
stuff2 <EOL>
await things2 ( x ) <EOL>
try : <EOL>
<INDENT>
STUFF2 <EOL>
<OUTDENT>
except MyException2 : <EOL>
<INDENT>
running2 = False <EOL>
<OUTDENT>
AFTER <EOL>
<OUTDENT>
try : <EOL>
<INDENT>
STUFF <EOL>
<OUTDENT>
except MyException : <EOL>
<INDENT>
running = False <EOL>
<OUTDENT>
BLOCK_AFTER <EOL>
<OUTDENT>
<OUTDENT>
except MyException2 : <EOL>
<INDENT>
OUTEREXCEPT <EOL>
<OUTDENT>
<OUTDENT>
print ( 'outside function' ) <EOL>
Loading

0 comments on commit 0540b39

Please sign in to comment.