Implement disabling imports in try/except and * imports, report errors on bad usage of lazy

Author: DinoV

Include/internal/pycore_magic_number.h

@@ -300,7 +300,7 @@ PC/launcher.c must also be updated.
 
 */
 
-#define PYC_MAGIC_NUMBER 3656
+#define PYC_MAGIC_NUMBER 3657
 /* This is equivalent to converting PYC_MAGIC_NUMBER to 2 bytes
    (little-endian) and then appending b'\r\n'. */
 #define PYC_MAGIC_NUMBER_TOKEN \

Lib/test/test_import/__init__.py

@@ -2662,6 +2662,50 @@ def test_lazy_value_get(self):
 
         self.assertTrue("test.test_import.data.lazy_imports.basic2" in sys.modules)
 
+    def test_lazy_try_except(self):
+        with self.assertRaises(SyntaxError):
+            import test.test_import.data.lazy_imports.lazy_try_except
+
+    def test_lazy_try_except_from(self):
+        with self.assertRaises(SyntaxError):
+            import test.test_import.data.lazy_imports.lazy_try_except_from
+
+    def test_lazy_try_except_from_star(self):
+        with self.assertRaises(SyntaxError):
+            import test.test_import.data.lazy_imports.lazy_try_except_from_star
+
+    def test_try_except_eager(self):
+        importlib.set_lazy_imports(True)
+        try:
+            import test.test_import.data.lazy_imports.try_except_eager
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertTrue("test.test_import.data.lazy_imports.basic2" in sys.modules)
+    
+    def test_try_except_eager_from(self):
+        importlib.set_lazy_imports(True)
+        try:
+            import test.test_import.data.lazy_imports.try_except_eager_from
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertTrue("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
+
+    def test_eager_import_func(self):
+        importlib.set_lazy_imports(True)
+        try:
+            import test.test_import.data.lazy_imports.eager_import_func
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        f = test.test_import.data.lazy_imports.eager_import_func.f
+        self.assertEqual(type(f()), type(sys))
+
 
 class TestSinglePhaseSnapshot(ModuleSnapshot):
     """A representation of a single-phase init module for testing.

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

@@ -0,0 +1,2 @@
+def f():
+    lazy import foo

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

@@ -0,0 +1,4 @@
+try:
+    lazy import foo
+except:
+    pass

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

@@ -0,0 +1,4 @@
+try:
+    lazy from foo import bar
+except:
+    pass

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

@@ -0,0 +1 @@
+lazy from foo import *

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

@@ -0,0 +1,4 @@
+try:
+    import test.test_import.data.lazy_imports.basic2
+except:
+    pass

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

@@ -0,0 +1,4 @@
+try:
+    from test.test_import.data.lazy_imports.basic2 import f
+except:
+    pass

Python/ceval.c

@@ -3049,17 +3049,6 @@ _PyEval_LazyImportName(PyThreadState *tstate, PyObject *builtins, PyObject *glob
             break;
     }
 
-    // Always make star imports eager regardless of lazy setting
-    if (fromlist && PyTuple_Check(fromlist) && PyTuple_GET_SIZE(fromlist) == 1) {
-        PyObject *item = PyTuple_GET_ITEM(fromlist, 0);
-        if (PyUnicode_Check(item)) {
-            const char *item_str = PyUnicode_AsUTF8(item);
-            if (item_str && strcmp(item_str, "*") == 0) {
-                lazy = 0;  // Force star imports to be eager
-            }
-        }
-    }
-
     if (!lazy) {
         // Not a lazy import or lazy imports are disabled, fallback to the regular import
         return _PyEval_ImportName(tstate, builtins, globals, locals, name, fromlist, level);

Python/codegen.c

@@ -2852,6 +2852,18 @@ codegen_import_as(compiler *c, location loc,
     return codegen_nameop(c, loc, asname, Store);
 }
 
+static int
+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;
+}
+
 static int
 codegen_import(compiler *c, stmt_ty s)
 {
@@ -2873,11 +2885,15 @@ codegen_import(compiler *c, stmt_ty s)
         ADDOP_LOAD_CONST(c, loc, zero);
         ADDOP_LOAD_CONST(c, loc, Py_None);
         if (s->v.Import.is_lazy) {
-            // TODO: SyntaxError when not in module scope
+            RETURN_IF_ERROR(codegen_validate_lazy_import(c, loc));
             ADDOP_NAME_CUSTOM(c, loc, IMPORT_NAME, alias->name, names, 2, 1);
         } else {
-            // TODO: If in try/except, set 2nd bit
-            ADDOP_NAME_CUSTOM(c, loc, IMPORT_NAME, alias->name, names, 2, 0);
+            if (_PyCompile_TopFBlock(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 {
+                ADDOP_NAME_CUSTOM(c, loc, IMPORT_NAME, alias->name, names, 2, 0);
+            }
         }
 
         if (alias->asname) {
@@ -2929,11 +2945,23 @@ codegen_from_import(compiler *c, stmt_ty s)
         from = s->v.ImportFrom.module;
     }
     if (s->v.ImportFrom.is_lazy) {
-        // TODO: SyntaxError when not in module scope
+        alias_ty alias = (alias_ty)asdl_seq_GET(s->v.ImportFrom.names, 0);
+        if (PyUnicode_READ_CHAR(alias->name, 0) == '*') {
+            return _PyCompile_Error(c, LOC(s), "cannot lazy import *");
+        }
+        RETURN_IF_ERROR(codegen_validate_lazy_import(c, LOC(s)));
         ADDOP_NAME_CUSTOM(c, LOC(s), IMPORT_NAME, from, names, 2, 1);
     } else {
-        ADDOP_NAME_CUSTOM(c, LOC(s), IMPORT_NAME, from, names, 2, 0);
+        alias_ty alias = (alias_ty)asdl_seq_GET(s->v.ImportFrom.names, 0);
+        if (_PyCompile_TopFBlock(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);
+        } else {
+            ADDOP_NAME_CUSTOM(c, LOC(s), IMPORT_NAME, from, names, 2, 0);
+        }
     }
+    
     for (Py_ssize_t i = 0; i < n; i++) {
         alias_ty alias = (alias_ty)asdl_seq_GET(s->v.ImportFrom.names, i);
         identifier store_name;