Allow try/except in with block

Author: DinoV

Include/internal/pycore_compile.h

@@ -127,6 +127,7 @@ int _PyCompile_PushFBlock(struct _PyCompiler *c, _Py_SourceLocation loc,
 void _PyCompile_PopFBlock(struct _PyCompiler *c, enum _PyCompile_FBlockType t,
                           _PyJumpTargetLabel block_label);
 _PyCompile_FBlockInfo *_PyCompile_TopFBlock(struct _PyCompiler *c);
+bool _PyCompile_InExceptionHandler(struct _PyCompiler *c);
 
 int _PyCompile_EnterScope(struct _PyCompiler *c, identifier name, int scope_type,
                           void *key, int lineno, PyObject *private,

Include/internal/pycore_symtable.h

@@ -127,7 +127,6 @@ typedef struct _symtable_entry {
     unsigned ste_has_conditional_annotations : 1; /* true if block has conditionally executed annotations */
     unsigned ste_in_conditional_block : 1; /* set while we are inside a conditionally executed block */
     unsigned ste_in_try_block : 1; /* set while we are inside a try/except block */
-    unsigned ste_in_with_block : 1; /* set while we are inside a with block */
     unsigned ste_in_unevaluated_annotation : 1; /* set while we are processing an annotation that will not be evaluated */
     int ste_comp_iter_expr; /* non-zero if visiting a comprehension range expression */
     _Py_SourceLocation ste_loc; /* source location of block */

Lib/test/test_import/__init__.py

@@ -2756,6 +2756,22 @@ def test_try_except_eager_from(self):
 
         self.assertTrue("test.test_import.data.lazy_imports.basic2" in sys.modules)
 
+    def test_lazy_with(self):
+        try:
+            import test.test_import.data.lazy_imports.lazy_with
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertFalse("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
+    def test_lazy_with_from(self):
+        try:
+            import test.test_import.data.lazy_imports.lazy_with_from
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertFalse("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
     def test_lazy_import_func(self):
         with self.assertRaises(SyntaxError):
             import test.test_import.data.lazy_imports.lazy_import_func

Lib/test/test_import/data/lazy_imports/lazy_with.py

@@ -0,0 +1,3 @@
+import contextlib
+with contextlib.nullcontext():
+    lazy import test.test_import.data.lazy_imports.basic2

Lib/test/test_import/data/lazy_imports/lazy_with_from.py

@@ -0,0 +1,3 @@
+import contextlib
+with contextlib.nullcontext():
+    lazy import test.test_import.data.lazy_imports.basic2 as basic2

Python/codegen.c

@@ -2857,8 +2857,6 @@ codegen_validate_lazy_import(compiler *c, location loc)
 {
     if (_PyCompile_ScopeType(c) != COMPILE_SCOPE_MODULE) {
         return _PyCompile_Error(c, loc, "lazy imports only allowed in module scope");
-    } else if (_PyCompile_TopFBlock(c)) {
-        return _PyCompile_Error(c, loc, "cannot lazy import in a nested scope");
     }
 
     return SUCCESS;
@@ -2888,7 +2886,7 @@ codegen_import(compiler *c, stmt_ty s)
             RETURN_IF_ERROR(codegen_validate_lazy_import(c, loc));
             ADDOP_NAME_CUSTOM(c, loc, IMPORT_NAME, alias->name, names, 2, 1);
         } else {
-            if (_PyCompile_TopFBlock(c) || _PyCompile_ScopeType(c) != COMPILE_SCOPE_MODULE) {
+            if (_PyCompile_InExceptionHandler(c) || _PyCompile_ScopeType(c) != COMPILE_SCOPE_MODULE) {
                 // force eager import in try/except block
                 ADDOP_NAME_CUSTOM(c, loc, IMPORT_NAME, alias->name, names, 2, 2);
             } else {
@@ -2953,7 +2951,8 @@ codegen_from_import(compiler *c, stmt_ty s)
         ADDOP_NAME_CUSTOM(c, LOC(s), IMPORT_NAME, from, names, 2, 1);
     } else {
         alias_ty alias = (alias_ty)asdl_seq_GET(s->v.ImportFrom.names, 0);
-        if (_PyCompile_TopFBlock(c) || _PyCompile_ScopeType(c) != COMPILE_SCOPE_MODULE ||
+        if (_PyCompile_InExceptionHandler(c) ||
+            _PyCompile_ScopeType(c) != COMPILE_SCOPE_MODULE ||
             PyUnicode_READ_CHAR(alias->name, 0) == '*') {
             // forced non-lazy import due to try/except or import *
             ADDOP_NAME_CUSTOM(c, LOC(s), IMPORT_NAME, from, names, 2, 2);

Python/compile.c

@@ -787,6 +787,26 @@ _PyCompile_TopFBlock(compiler *c)
     return &c->u->u_fblock[c->u->u_nfblocks - 1];
 }
 
+bool
+_PyCompile_InExceptionHandler(compiler *c)
+{
+    for (Py_ssize_t i = c->u->u_nfblocks; i < c->u->u_nfblocks; i++) {
+        fblockinfo *block = &c->u->u_fblock[i];
+        switch (block->fb_type) {
+            case COMPILE_FBLOCK_TRY_EXCEPT:
+            case COMPILE_FBLOCK_FINALLY_TRY:
+            case COMPILE_FBLOCK_FINALLY_END:
+            case COMPILE_FBLOCK_EXCEPTION_HANDLER:
+            case COMPILE_FBLOCK_EXCEPTION_GROUP_HANDLER:
+            case COMPILE_FBLOCK_HANDLER_CLEANUP:
+                return true;
+            default:
+                break;
+        }
+    }    
+    return false;
+}
+
 void
 _PyCompile_DeferredAnnotations(compiler *c,
                                PyObject **deferred_annotations,

Python/symtable.c

@@ -142,7 +142,6 @@ ste_new(struct symtable *st, identifier name, _Py_block_ty block,
     ste->ste_has_conditional_annotations = 0;
     ste->ste_in_conditional_block = 0;
     ste->ste_in_try_block = 0;
-    ste->ste_in_with_block = 0;
     ste->ste_in_unevaluated_annotation = 0;
     ste->ste_annotation_block = NULL;
 
@@ -1756,13 +1755,6 @@ symtable_enter_type_param_block(struct symtable *st, identifier name,
 #define LEAVE_TRY_BLOCK(ST) \
     (ST)->st_cur->ste_in_try_block = in_try_block;
 
-#define ENTER_WITH_BLOCK(ST) \
-    int in_with_block = (ST)->st_cur->ste_in_with_block; \
-    (ST)->st_cur->ste_in_with_block = 1;
-
-#define LEAVE_WITH_BLOCK(ST) \
-    (ST)->st_cur->ste_in_with_block = in_with_block;
-
 #define ENTER_RECURSIVE() \
 if (Py_EnterRecursiveCall(" during compilation")) { \
     return 0; \
@@ -1835,14 +1827,6 @@ check_lazy_import_context(struct symtable *st, stmt_ty s, const char* import_typ
         return 0;
     }
 
-    /* Check if inside with block */
-    if (st->st_cur->ste_in_with_block) {
-        PyErr_Format(PyExc_SyntaxError,
-                    "lazy %s not allowed inside with blocks", import_type);
-        SET_ERROR_LOCATION(st->st_filename, LOCATION(s));
-        return 0;
-    }
-
     /* Check if inside function scope */
     if (st->st_cur->ste_type == FunctionBlock) {
         PyErr_Format(PyExc_SyntaxError,
@@ -2258,10 +2242,8 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
         break;
     case With_kind: {
         ENTER_CONDITIONAL_BLOCK(st);
-        ENTER_WITH_BLOCK(st);
         VISIT_SEQ(st, withitem, s->v.With.items);
         VISIT_SEQ(st, stmt, s->v.With.body);
-        LEAVE_WITH_BLOCK(st);
         LEAVE_CONDITIONAL_BLOCK(st);
         break;
     }
@@ -2326,10 +2308,8 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
             return 0;
         }
         ENTER_CONDITIONAL_BLOCK(st);
-        ENTER_WITH_BLOCK(st);
         VISIT_SEQ(st, withitem, s->v.AsyncWith.items);
         VISIT_SEQ(st, stmt, s->v.AsyncWith.body);
-        LEAVE_WITH_BLOCK(st);
         LEAVE_CONDITIONAL_BLOCK(st);
         break;
     }