Implement Python 3.8 && 3.9 try-except-finally logic

ASTree.cpp

@@ -90,7 +90,6 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
     int pos = 0;
     int unpack = 0;
     bool else_pop = false;
-    bool need_try = false;
     bool variable_annotations = false;
 
     while (!source.atEof()) {
@@ -110,15 +109,21 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
         curpos = pos;
         bc_next(source, mod, opcode, operand, pos);
 
-        if (need_try && opcode != Pyc::SETUP_EXCEPT_A) {
-            need_try = false;
+        if (mod->verCompare(3, 9) >= 0
+            && opcode == Pyc::JUMP_FORWARD_A
+            && curblock->blktype() == ASTBlock::BLK_FINALLY)
+        {
+            /* Ignore all opcodes contained between this jump until
+               its target when we are dealing with Python 3.9 and 3.10
+               as they duplicate the code related to the FINALLY
+               block. Lukily, the target of this JUMP_FORWARD_A
+               exactly surrounds the duplicated code */
+            pos = pos + operand;
+            source.setPos(pos);
+            opcode = Pyc::END_FINALLY;
+        }
 
-            /* Store the current stack for the except/finally statement(s) */
-            stack_hist.push(stack);
-            PycRef<ASTBlock> tryblock = new ASTBlock(ASTBlock::BLK_TRY, curblock->end(), true);
-            blocks.push(tryblock);
-            curblock = blocks.top();
-        } else if (else_pop
+        if (else_pop
                 && opcode != Pyc::JUMP_FORWARD_A
                 && opcode != Pyc::JUMP_IF_FALSE_A
                 && opcode != Pyc::JUMP_IF_FALSE_OR_POP_A
@@ -128,6 +133,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                 && opcode != Pyc::JUMP_IF_TRUE_OR_POP_A
                 && opcode != Pyc::POP_JUMP_IF_TRUE_A
                 && opcode != Pyc::POP_JUMP_FORWARD_IF_TRUE_A
+                && opcode != Pyc::JUMP_IF_NOT_EXC_MATCH_A
                 && opcode != Pyc::POP_BLOCK) {
             else_pop = false;
 
@@ -802,6 +808,8 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                 }
             }
             break;
+        case Pyc::RERAISE:
+        case Pyc::RERAISE_A:
         case Pyc::END_FINALLY:
             {
                 bool isFinally = false;
@@ -1055,11 +1063,25 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
         case Pyc::POP_JUMP_FORWARD_IF_TRUE_A:
         case Pyc::INSTRUMENTED_POP_JUMP_IF_FALSE_A:
         case Pyc::INSTRUMENTED_POP_JUMP_IF_TRUE_A:
+        case Pyc::JUMP_IF_NOT_EXC_MATCH_A:
             {
                 PycRef<ASTNode> cond = stack.top();
                 PycRef<ASTCondBlock> ifblk;
                 int popped = ASTCondBlock::UNINITED;
 
+                /* Emulate a CMP_EXCEPTION condition block before
+                   evaluating the jump */
+                if (opcode == Pyc::JUMP_IF_NOT_EXC_MATCH_A)
+                {
+                    PycRef<ASTNode> exc_type = stack.top();
+                    stack.pop();
+                    PycRef<ASTNode> raised_exc = stack.top();
+                    stack.pop();
+
+                    cond = new ASTCompare(raised_exc, exc_type, ASTCompare::CMP_EXCEPTION);
+                    popped = ASTCondBlock::PRE_POPPED;
+                }
+
                 if (opcode == Pyc::POP_JUMP_IF_FALSE_A
                         || opcode == Pyc::POP_JUMP_IF_TRUE_A
                         || opcode == Pyc::POP_JUMP_FORWARD_IF_FALSE_A
@@ -1589,6 +1611,18 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
             break;
         case Pyc::POP_BLOCK:
             {
+                if (mod->verCompare(3,9)>=0)
+                {
+                    /* Emulate a BEGIN_FINALLY for Python 3.9 when
+                       POP_BLOCK happens in a FINALLY block */
+                    int target_opcode, target_operand;
+                    bc_get_at(source, mod, target_opcode, target_operand, pos);
+                    if (target_opcode != Pyc::JUMP_FORWARD_A)
+                    {
+                        stack.push(NULL);
+                    }
+                }
+
                 if (curblock->blktype() == ASTBlock::BLK_CONTAINER ||
                         curblock->blktype() == ASTBlock::BLK_FINALLY) {
                     /* These should only be popped by an END_FINALLY */
@@ -1918,31 +1952,92 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
         case Pyc::WITH_CLEANUP_FINISH:
             /* Ignore this */
             break;
-        case Pyc::SETUP_EXCEPT_A:
+        case Pyc::BEGIN_FINALLY:
             {
-                if (curblock->blktype() == ASTBlock::BLK_CONTAINER) {
-                    curblock.cast<ASTContainerBlock>()->setExcept(pos+operand);
-                } else {
-                    PycRef<ASTBlock> next = new ASTContainerBlock(0, pos+operand);
-                    blocks.push(next.cast<ASTBlock>());
-                }
-
-                /* Store the current stack for the except/finally statement(s) */
-                stack_hist.push(stack);
-                PycRef<ASTBlock> tryblock = new ASTBlock(ASTBlock::BLK_TRY, pos+operand, true);
-                blocks.push(tryblock.cast<ASTBlock>());
-                curblock = blocks.top();
-
-                need_try = false;
+                /* Only in Python 3.8: Push NULL on the stack */
+                stack.push(NULL);
             }
             break;
+        case Pyc::SETUP_EXCEPT_A:
         case Pyc::SETUP_FINALLY_A:
             {
-                PycRef<ASTBlock> next = new ASTContainerBlock(pos+operand);
-                blocks.push(next.cast<ASTBlock>());
-                curblock = blocks.top();
+                /* Up til Python 3.7 */
+                bool is_except = opcode == Pyc::SETUP_EXCEPT_A;
 
-                need_try = true;
+                /* Starting from Python 3.8 */
+                if (mod->verCompare(3, 8) >= 0)
+                {
+                    /* Try to guess if we have an exception or finally
+                       block by looking at the opcode pointed by the operand */
+                    int target_opcode, target_operand;
+                    bc_get_at(source, mod, target_opcode, target_operand, pos+operand);
+
+                    /* POP_TOP (x3) and DUP_TOP usually mean it's the start of
+                       an exception block */
+                    if (target_opcode == Pyc::POP_TOP || target_opcode == Pyc::DUP_TOP)
+                    {
+                      is_except = true;
+                    }
+                }
+
+                if (is_except)
+                {
+                    if (curblock->blktype() == ASTBlock::BLK_CONTAINER)
+                    {
+                        curblock.cast<ASTContainerBlock>()->setExcept(pos+operand);
+                    }
+                    else
+                    {
+                        PycRef<ASTBlock> next = new ASTContainerBlock(0, pos+operand);
+                        blocks.push(next.cast<ASTBlock>());
+                    }
+
+                    /* Store the current stack for the except/finally statement(s) */
+                    stack_hist.push(stack);
+                    PycRef<ASTBlock> tryblock = new ASTBlock(ASTBlock::BLK_TRY, pos+operand, true);
+                    blocks.push(tryblock.cast<ASTBlock>());
+                    curblock = blocks.top();
+                }
+                /* else it's a finally block */
+                else
+                {
+                    PycRef<ASTBlock> next = new ASTContainerBlock(pos+operand);
+                    blocks.push(next.cast<ASTBlock>());
+                    curblock = blocks.top();
+
+                    /* Look at the next opcode. If it is not an EXCEPT
+                       block, then push a new TRY block */
+                    int next_opcode, next_operand;
+                    int next_pos = bc_get_at(source, mod, next_opcode, next_operand, pos);
+
+                    /* Up til Python 3.7 */
+                    bool is_next_except = next_opcode == Pyc::SETUP_EXCEPT_A;
+
+                    /* Starting from Python 3.8 */
+                    if (mod->verCompare(3, 8) >= 0 && next_opcode == Pyc::SETUP_FINALLY_A)
+                    {
+                        /* Try to guess if we have an exception or finally
+                           block by looking at the opcode pointed by the operand */
+                        int target_opcode, target_operand;
+                        bc_get_at(source, mod, target_opcode, target_operand, next_pos+next_operand);
+
+                        /* POP_TOP (x3) and DUP_TOP usually mean it's the start of
+                           an exception block */
+                        if (target_opcode == Pyc::POP_TOP || target_opcode == Pyc::DUP_TOP)
+                        {
+                            is_next_except = true;
+                        }
+                    }
+
+                    if ( !is_next_except )
+                    {
+                        /* Store the current stack for the except/finally statement(s) */
+                        stack_hist.push(stack);
+                        PycRef<ASTBlock> tryblock = new ASTBlock(ASTBlock::BLK_TRY, curblock->end(), true);
+                        blocks.push(tryblock);
+                        curblock = blocks.top();
+                    }
+                }
             }
             break;
         case Pyc::SETUP_LOOP_A:

CMakeLists.txt

@@ -16,6 +16,10 @@ if (ENABLE_STACK_DEBUG)
     add_definitions(-DSTACK_DEBUG)
 endif()
 
+#if (ENABLE_BLOCK_DEBUG OR ENABLE_STACK_DEBUG)
+    set(CMAKE_CXX_FLAGS "-g ${CMAKE_CXX_FLAGS}")
+#endif()
+
 if(CMAKE_COMPILER_IS_GNUCXX OR "${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang")
     set(CMAKE_CXX_FLAGS "-Wall -Wextra -Wno-error=shadow -Werror ${CMAKE_CXX_FLAGS}")
 elseif(MSVC)

bytecode.cpp

@@ -300,6 +300,25 @@ void bc_next(PycBuffer& source, PycModule* mod, int& opcode, int& operand, int&
     }
 }
 
+int bc_get_at(PycBuffer& source, PycModule *mod, int& opcode, int &operand, int pos)
+{
+    /* save current pos */
+    int old_pos = source.setPos(pos);
+    if (old_pos == EOF)
+    {
+        opcode = 0;
+        operand = 0;
+        return EOF;
+    }
+
+    /* read opcode at requested pos */
+    bc_next(source, mod, opcode, operand, pos);
+
+    /* restore saved pos */
+    source.setPos(old_pos);
+    return pos;
+}
+
 void bc_disasm(std::ostream& pyc_output, PycRef<PycCode> code, PycModule* mod,
                int indent, unsigned flags)
 {

bytecode.h

@@ -30,5 +30,6 @@ int ByteToOpcode(int maj, int min, int opcode);
 void print_const(std::ostream& pyc_output, PycRef<PycObject> obj, PycModule* mod,
                  const char* parent_f_string_quote = nullptr);
 void bc_next(PycBuffer& source, PycModule* mod, int& opcode, int& operand, int& pos);
+int bc_get_at(PycBuffer& source, PycModule *mod, int& opcode, int &operand, int pos);
 void bc_disasm(std::ostream& pyc_output, PycRef<PycCode> code, PycModule* mod,
                int indent, unsigned flags);

data.cpp

@@ -63,6 +63,17 @@ int PycFile::getBuffer(int bytes, void* buffer)
     return (int)fread(buffer, 1, bytes, m_stream);
 }
 
+int PycFile::setPos(int pos)
+{
+    int old_pos = ftell(m_stream);
+    if (fseek(m_stream, pos, SEEK_SET) != 0)
+    {
+        fseek(m_stream, old_pos, SEEK_SET);
+        return EOF;
+    }
+    return old_pos;
+}
+
 
 /* PycBuffer */
 int PycBuffer::getByte()
@@ -74,6 +85,13 @@ int PycBuffer::getByte()
     return ch & 0xFF;   // Make sure it's just a byte!
 }
 
+int PycBuffer::setPos(int pos)
+{
+    int old_pos = m_pos;
+    m_pos = pos;
+    return old_pos;
+}
+
 int PycBuffer::getBuffer(int bytes, void* buffer)
 {
     if (m_pos + bytes > m_size)

data.h

@@ -20,6 +20,7 @@ class PycData {
 
     virtual int getByte() = 0;
     virtual int getBuffer(int bytes, void* buffer) = 0;
+    virtual int setPos(int pos) = 0;
     int get16();
     int get32();
     Pyc_INT64 get64();
@@ -35,6 +36,7 @@ class PycFile : public PycData {
 
     int getByte() override;
     int getBuffer(int bytes, void* buffer) override;
+    int setPos(int pos) override;
 
 private:
     FILE* m_stream;
@@ -51,6 +53,7 @@ class PycBuffer : public PycData {
 
     int getByte() override;
     int getBuffer(int bytes, void* buffer) override;
+    int setPos(int pos) override;
 
 private:
     const unsigned char* m_buffer;

tests/input/try_except.py

@@ -0,0 +1,16 @@
+try:
+    import sys
+    try:
+        print ('something else')
+    except AssertionError:
+        print ('...failed')
+except ImportError:
+    print ('Oh Noes!')
+
+try:
+    try:
+        print("try")
+    except:
+        print("except in")
+except:
+    print("except out")

tests/input/try_except_finally.py

@@ -1,10 +1,10 @@
 try:
     import sys
     try:
-        print 'something else'
+        print ('something else')
     except AssertionError:
-        print '...failed'
+        print ('...failed')
 except ImportError:
-    print 'Oh Noes!'
+    print ('Oh Noes!')
 finally:
-    print 'Exiting'
+    print ('Exiting')

tests/input/try_except_finally_py2.py

@@ -0,0 +1,10 @@
+try:
+    import sys
+    try:
+        print 'something else'
+    except AssertionError:
+        print '...failed'
+except ImportError:
+    print 'Oh Noes!'
+finally:
+    print 'Exiting'

tests/input/try_finally.py

@@ -0,0 +1,17 @@
+try:
+    import sys
+    try:
+        print ('something else')
+    finally:
+        print ('...ok')
+finally:
+    print ('Exiting')
+
+
+try:
+    try:
+        print("try")
+    finally:
+        print("finally in")
+finally:
+    print("finally out")

tests/input/try_finally_simple.py

@@ -0,0 +1,7 @@
+try:
+    try:
+        print("try")
+    finally:
+        print("finally in")
+finally:
+    print("finally out")