depyo 1.0.2 → 1.0.3

package/lib/handlers/context_managers.js

@@ -33,7 +33,10 @@ function extractWithPattern() {
 
     if (thirdInstr && thirdInstr.OpCodeID === this.OpCodes.STORE_FAST_A) {
         // Has "as variable"
-        asVariable = new AST.ASTName(thirdInstr.Name?.toString() || "###FIXME###");
+        if (thirdInstr.Name == null) {
+            throw new Error(`SETUP_WITH at offset ${this.code.Current.Offset} (file offset ${this.object.codeOffset + this.code.Current.Offset}): STORE_FAST has null Name (expected "as" variable)`);
+        }
+        asVariable = new AST.ASTName(thirdInstr.Name.toString());
     } else if (thirdInstr && thirdInstr.OpCodeID === this.OpCodes.POP_TOP) {
         // No "as" clause
         asVariable = null;
@@ -62,11 +65,14 @@ function extractWithPattern() {
                     instr.OpCodeID === this.OpCodes.LOAD_GLOBAL_A ||
                     instr.OpCodeID === this.OpCodes.LOAD_NAME_A ||
                     instr.OpCodeID === this.OpCodes.LOAD_DEREF_A) {
-                    ctxMgrExpr = new AST.ASTName(instr.Name?.toString() || "###FIXME###");
-                }else {
-                    // For CALL_FUNCTION, LOAD_ATTR, etc., we can't easily reconstruct
-                    // the expression without replaying the stack. Use placeholder for now.
-                    ctxMgrExpr = new AST.ASTName("###FIXME###");
+                    if (instr.Name == null) {
+                        throw new Error(`SETUP_WITH at offset ${this.code.Current.Offset} (file offset ${this.object.codeOffset + this.code.Current.Offset}): ${instr.InstructionName} at offset ${instr.Offset} has null Name`);
+                    }
+                    ctxMgrExpr = new AST.ASTName(instr.Name.toString());
+                } else {
+                    // CALL_FUNCTION/LOAD_ATTR/BINARY_SUBSCR: reconstructing the
+                    // expression without replaying the stack isn't supported yet.
+                    throw new Error(`SETUP_WITH at offset ${this.code.Current.Offset} (file offset ${this.object.codeOffset + this.code.Current.Offset}): unsupported context manager opcode ${instr.InstructionName} at offset ${instr.Offset} — needs stack replay`);
                 }
                 break;
             }
@@ -117,6 +123,102 @@ function handleWithCleanupStart() {
 }
 
 function handleWithCleanup() {
+    // Async-with closure runs first because the 3.8+ pattern uses BEGIN_FINALLY
+    // (no-op handler) instead of LOAD_CONST None — so the data stack may be
+    // empty when WITH_CLEANUP_START fires. Don't gate on the None pop here.
+    if (this.curBlock.blockType == AST.ASTBlock.BlockType.AsyncWith
+            && this.curBlock.end == this.code.Current.Offset) {
+        let asyncWithBlock = this.curBlock;
+
+        // Skip the standard 4-op cleanup epilogue:
+        //   GET_AWAITABLE, LOAD_CONST None, YIELD_FROM, WITH_CLEANUP_FINISH
+        let standardEpilogue = [
+            this.OpCodes.GET_AWAITABLE,
+            this.OpCodes.LOAD_CONST_A,
+            this.OpCodes.YIELD_FROM,
+            this.OpCodes.WITH_CLEANUP_FINISH,
+        ];
+        for (let expectedOp of standardEpilogue) {
+            if (this.code.Next?.OpCodeID === expectedOp) {
+                this.code.GoNext();
+            } else {
+                break;
+            }
+        }
+
+        // After the standard epilogue, two patterns are possible (3.8+):
+        //   Normal exit:        END_FINALLY
+        //   Early return path:  POP_FINALLY 0 / LOAD_CONST X / RETURN_VALUE
+        //                       / WITH_CLEANUP_START / GET_AWAITABLE / LOAD_CONST None
+        //                       / YIELD_FROM / WITH_CLEANUP_FINISH / END_FINALLY
+        //                       [/ LOAD_CONST None / RETURN_VALUE  -- dead implicit return]
+        // For the early-return case we synthesize the explicit return into the body
+        // and skip the trailing exception handler + dead implicit return.
+        let after = this.code.Next;
+        if (after?.OpCodeID === this.OpCodes.POP_FINALLY_A
+                || after?.OpCodeID === this.OpCodes.POP_FINALLY) {
+            this.code.GoNext();
+
+            let loadConst = this.code.Next;
+            let returnValue = this.code.PeekNextInstruction(2);
+            if (loadConst?.OpCodeID === this.OpCodes.LOAD_CONST_A
+                    && returnValue?.OpCodeID === this.OpCodes.RETURN_VALUE) {
+                let constObj = new AST.ASTObject(loadConst.ConstantObject);
+                constObj.line = loadConst.LineNo;
+                let retValue;
+                if (constObj.object == null || constObj.object.ClassName == "Py_None") {
+                    retValue = new AST.ASTNone();
+                } else {
+                    retValue = constObj;
+                }
+                let retNode = new AST.ASTReturn(retValue);
+                retNode.line = loadConst.LineNo;
+                asyncWithBlock.append(retNode);
+                this.code.GoNext();
+                this.code.GoNext();
+            }
+
+            let exceptionHandler = [
+                this.OpCodes.WITH_CLEANUP_START,
+                this.OpCodes.GET_AWAITABLE,
+                this.OpCodes.LOAD_CONST_A,
+                this.OpCodes.YIELD_FROM,
+                this.OpCodes.WITH_CLEANUP_FINISH,
+                this.OpCodes.END_FINALLY,
+            ];
+            for (let expectedOp of exceptionHandler) {
+                if (this.code.Next?.OpCodeID === expectedOp) {
+                    this.code.GoNext();
+                } else {
+                    break;
+                }
+            }
+
+            // Skip the dead implicit LOAD_CONST None / RETURN_VALUE if it's
+            // the very tail of the function (no more instructions after).
+            let load = this.code.Next;
+            let ret = this.code.PeekNextInstruction(2);
+            let beyond = this.code.PeekNextInstruction(3);
+            if (load?.OpCodeID === this.OpCodes.LOAD_CONST_A
+                    && ret?.OpCodeID === this.OpCodes.RETURN_VALUE
+                    && beyond == null) {
+                this.code.GoNext();
+                this.code.GoNext();
+            }
+        } else if (after?.OpCodeID === this.OpCodes.END_FINALLY) {
+            this.code.GoNext();
+        }
+
+        this.blocks.pop();
+        this.curBlock = this.blocks.top();
+        this.curBlock.append(asyncWithBlock);
+
+        if (global.g_cliArgs?.debug) {
+            console.log(`ASYNC WITH block closed, nodes count=${asyncWithBlock.nodes.length}`);
+        }
+        return;
+    }
+
     // Stack top should be a None. Ignore it.
     let none = this.dataStack.pop();
 
@@ -226,19 +328,106 @@ function handleBeforeWith() {
     }
 }
 
+/**
+ * Handler for BEFORE_ASYNC_WITH (Python 3.5-3.10).
+ *
+ * Bytecode pattern for `async with EXPR [as VAR]:`:
+ *   LOAD_*               EXPR (context manager)
+ *   BEFORE_ASYNC_WITH    <-- we are here
+ *   GET_AWAITABLE
+ *   LOAD_CONST None
+ *   YIELD_FROM
+ *   SETUP_ASYNC_WITH +N  (N targets WITH_CLEANUP_START — end of body)
+ *   POP_TOP | STORE_FAST VAR
+ *   ... body ...
+ *   POP_BLOCK
+ *   LOAD_CONST None
+ *   WITH_CLEANUP_START   (closes block, see handleWithCleanup)
+ *   GET_AWAITABLE
+ *   LOAD_CONST None
+ *   YIELD_FROM
+ *   WITH_CLEANUP_FINISH
+ *   END_FINALLY
+ *
+ * We pop the context manager off the data stack, then use GoNext to consume
+ * the entire setup epilogue. This bypasses the GET_AWAITABLE / YIELD_FROM
+ * handlers (which would otherwise emit `await await(EXPR.__aenter__)` noise)
+ * and creates a clean ASTAsyncWithBlock for the body.
+ */
 function handleBeforeAsyncWith() {
-    let ctxmgr = this.dataStack.top();
-    let callNode = new AST.ASTCall(new AST.ASTName('await'), [new AST.ASTBinary(ctxmgr, new AST.ASTName('__aenter__'), AST.ASTBinary.BinOp.Attr)], []);
-    callNode.line = this.code.Current.LineNo;
-    this.dataStack.push(callNode);
-}
+    let ctxMgr = this.dataStack.pop();
 
-function handleSetupAsyncWithA() {
-    let asyncWithBlock = new AST.ASTBlock(AST.ASTBlock.BlockType.AsyncWith, this.code.Current.Offset, this.code.Current.JumpTarget);
+    let getAwaitable = this.code.PeekNextInstruction(1);
+    let loadNone = this.code.PeekNextInstruction(2);
+    let yieldFrom = this.code.PeekNextInstruction(3);
+    let setupAsyncWith = this.code.PeekNextInstruction(4);
+
+    if (!getAwaitable || getAwaitable.OpCodeID !== this.OpCodes.GET_AWAITABLE
+        || !loadNone || loadNone.OpCodeID !== this.OpCodes.LOAD_CONST_A
+        || !yieldFrom || yieldFrom.OpCodeID !== this.OpCodes.YIELD_FROM
+        || !setupAsyncWith || setupAsyncWith.OpCodeID !== this.OpCodes.SETUP_ASYNC_WITH_A) {
+        throw new Error(`BEFORE_ASYNC_WITH at offset ${this.code.Current.Offset}: expected GET_AWAITABLE/LOAD_CONST/YIELD_FROM/SETUP_ASYNC_WITH sequence, got ${getAwaitable?.InstructionName}/${loadNone?.InstructionName}/${yieldFrom?.InstructionName}/${setupAsyncWith?.InstructionName}`);
+    }
+
+    let asClause = this.code.PeekNextInstruction(5);
+    let asVar = null;
+    let extraSkip = 0;
+    if (asClause?.OpCodeID === this.OpCodes.STORE_FAST_A) {
+        if (asClause.Name == null) {
+            throw new Error(`BEFORE_ASYNC_WITH at offset ${this.code.Current.Offset}: STORE_FAST has null Name (expected "as" variable)`);
+        }
+        asVar = new AST.ASTName(asClause.Name.toString());
+        extraSkip = 1;
+    } else if (asClause?.OpCodeID === this.OpCodes.POP_TOP) {
+        extraSkip = 1;
+    }
+
+    // Walk forward from the body start to find the matching WITH_CLEANUP_START
+    // for THIS async-with (not a nested one). Track depth so nested
+    // SETUP_WITH/SETUP_ASYNC_WITH are skipped over correctly.
+    //
+    // Note: SETUP_ASYNC_WITH.JumpTarget is unreliable here — for the early-return
+    // pattern it points to the exception-path cleanup, not the normal exit.
+    let blockEnd = setupAsyncWith.JumpTarget;
+    let depth = 0;
+    let searchIdx = this.code.CurrentInstructionIndex + 5 + extraSkip;
+    for (let i = searchIdx; i < this.code.Instructions.length; i++) {
+        let instr = this.code.Instructions[i];
+        if (!instr) break;
+        if (instr.OpCodeID === this.OpCodes.SETUP_ASYNC_WITH_A
+            || instr.OpCodeID === this.OpCodes.SETUP_WITH_A) {
+            depth++;
+            continue;
+        }
+        if (instr.OpCodeID === this.OpCodes.WITH_CLEANUP_START
+            || instr.OpCodeID === this.OpCodes.WITH_CLEANUP) {
+            if (depth === 0) {
+                blockEnd = instr.Offset;
+                break;
+            }
+            depth--;
+        }
+    }
+
+    // Skip GET_AWAITABLE, LOAD_CONST None, YIELD_FROM, SETUP_ASYNC_WITH, [POP_TOP|STORE_FAST]
+    this.code.GoNext(4 + extraSkip);
+
+    let asyncWithBlock = new AST.ASTAsyncWithBlock(this.code.Current.Offset, blockEnd);
+    asyncWithBlock.expr = ctxMgr;
+    asyncWithBlock.var = asVar;
+    asyncWithBlock.init();
     this.blocks.push(asyncWithBlock);
     this.curBlock = this.blocks.top();
 }
 
+function handleSetupAsyncWithA() {
+    // Normally consumed by handleBeforeAsyncWith via GoNext. If we land here,
+    // the lookahead failed — bail silently rather than pushing a malformed block.
+    if (global.g_cliArgs?.debug) {
+        console.error(`SETUP_ASYNC_WITH at offset ${this.code.Current.Offset}: reached without preceding BEFORE_ASYNC_WITH consumption — ignoring`);
+    }
+}
+
 module.exports = {
     handleBeforeWith,
     handleBeforeAsyncWith,