Syntax restrictions for lazy imports

Author: pablogsal

Include/internal/pycore_symtable.h

@@ -126,6 +126,7 @@ typedef struct _symtable_entry {
     unsigned ste_method : 1; /* true if block is a function block defined in class scope */
     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_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_syntax.py

@@ -3394,6 +3394,119 @@ def test_ifexp_body_stmt_else_stmt(self):
         ]:
             self._check_error(f"x = {lhs_stmt} if 1 else {rhs_stmt}", msg)
 
+
+class LazyImportRestrictionTestCase(SyntaxErrorTestCase):
+    """Test syntax restrictions for lazy imports."""
+
+    def test_lazy_import_in_try_block(self):
+        """Test that lazy imports are not allowed inside try blocks."""
+        self._check_error("""\
+try:
+    lazy import os
+except:
+    pass
+""", "lazy import not allowed inside try/except blocks")
+
+        self._check_error("""\
+try:
+    lazy from sys import path
+except ImportError:
+    pass
+""", "lazy from ... import not allowed inside try/except blocks")
+
+    def test_lazy_import_in_trystar_block(self):
+        """Test that lazy imports are not allowed inside try* blocks."""
+        self._check_error("""\
+try:
+    lazy import json
+except* Exception:
+    pass
+""", "lazy import not allowed inside try/except blocks")
+
+        self._check_error("""\
+try:
+    lazy from collections import defaultdict
+except* ImportError:
+    pass
+""", "lazy from ... import not allowed inside try/except blocks")
+
+    def test_lazy_import_in_function(self):
+        """Test that lazy imports are not allowed inside functions."""
+        self._check_error("""\
+def func():
+    lazy import math
+""", "lazy import not allowed inside functions")
+
+        self._check_error("""\
+def func():
+    lazy from datetime import datetime
+""", "lazy from ... import not allowed inside functions")
+
+    def test_lazy_import_in_async_function(self):
+        """Test that lazy imports are not allowed inside async functions."""
+        self._check_error("""\
+async def async_func():
+    lazy import asyncio
+""", "lazy import not allowed inside functions")
+
+        self._check_error("""\
+async def async_func():
+    lazy from json import loads
+""", "lazy from ... import not allowed inside functions")
+
+    def test_lazy_import_in_class(self):
+        """Test that lazy imports are not allowed inside classes."""
+        self._check_error("""\
+class MyClass:
+    lazy import typing
+""", "lazy import not allowed inside classes")
+
+        self._check_error("""\
+class MyClass:
+    lazy from abc import ABC
+""", "lazy from ... import not allowed inside classes")
+
+    def test_lazy_import_star_forbidden(self):
+        """Test that 'lazy from ... import *' is forbidden everywhere."""
+        # At module level should also be forbidden
+        self._check_error("lazy from os import *",
+                         "lazy from ... import \\* is not allowed")
+
+        # Inside function should give lazy function error first
+        self._check_error("""\
+def func():
+    lazy from sys import *
+""", "lazy from ... import not allowed inside functions")
+
+    def test_lazy_import_nested_scopes(self):
+        """Test lazy imports in nested scopes."""
+        self._check_error("""\
+class Outer:
+    def method(self):
+        lazy import sys
+""", "lazy import not allowed inside functions")
+
+        self._check_error("""\
+def outer():
+    class Inner:
+        lazy import json
+""", "lazy import not allowed inside classes")
+
+        self._check_error("""\
+def outer():
+    def inner():
+        lazy from collections import deque
+""", "lazy from ... import not allowed inside functions")
+
+    def test_lazy_import_valid_cases(self):
+        """Test that lazy imports work at module level."""
+        # These should compile without errors
+        compile("lazy import os", "<test>", "exec")
+        compile("lazy from sys import path", "<test>", "exec")
+        compile("lazy import json as j", "<test>", "exec")
+        compile("lazy from datetime import datetime as dt", "<test>", "exec")
+
+
 def load_tests(loader, tests, pattern):
     tests.addTest(doctest.DocTestSuite())
     return tests

Python/symtable.c

@@ -1747,6 +1747,13 @@ symtable_enter_type_param_block(struct symtable *st, identifier name,
 #define LEAVE_CONDITIONAL_BLOCK(ST) \
     (ST)->st_cur->ste_in_conditional_block = in_conditional_block;
 
+#define ENTER_TRY_BLOCK(ST) \
+    int in_try_block = (ST)->st_cur->ste_in_try_block; \
+    (ST)->st_cur->ste_in_try_block = 1;
+
+#define LEAVE_TRY_BLOCK(ST) \
+    (ST)->st_cur->ste_in_try_block = in_try_block;
+
 #define ENTER_RECURSIVE() \
 if (Py_EnterRecursiveCall(" during compilation")) { \
     return 0; \
@@ -1808,6 +1815,36 @@ check_import_from(struct symtable *st, stmt_ty s)
     return 1;
 }
 
+static int
+check_lazy_import_context(struct symtable *st, stmt_ty s, const char* import_type)
+{
+    /* Check if inside try/except block */
+    if (st->st_cur->ste_in_try_block) {
+        PyErr_Format(PyExc_SyntaxError,
+                    "lazy %s not allowed inside try/except 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,
+                    "lazy %s not allowed inside functions", import_type);
+        SET_ERROR_LOCATION(st->st_filename, LOCATION(s));
+        return 0;
+    }
+
+    /* Check if inside class scope */
+    if (st->st_cur->ste_type == ClassBlock) {
+        PyErr_Format(PyExc_SyntaxError,
+                    "lazy %s not allowed inside classes", import_type);
+        SET_ERROR_LOCATION(st->st_filename, LOCATION(s));
+        return 0;
+    }
+
+    return 1;
+}
+
 static bool
 allows_top_level_await(struct symtable *st)
 {
@@ -2076,19 +2113,23 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
         break;
     case Try_kind: {
         ENTER_CONDITIONAL_BLOCK(st);
+        ENTER_TRY_BLOCK(st);
         VISIT_SEQ(st, stmt, s->v.Try.body);
         VISIT_SEQ(st, excepthandler, s->v.Try.handlers);
         VISIT_SEQ(st, stmt, s->v.Try.orelse);
         VISIT_SEQ(st, stmt, s->v.Try.finalbody);
+        LEAVE_TRY_BLOCK(st);
         LEAVE_CONDITIONAL_BLOCK(st);
         break;
     }
     case TryStar_kind: {
         ENTER_CONDITIONAL_BLOCK(st);
+        ENTER_TRY_BLOCK(st);
         VISIT_SEQ(st, stmt, s->v.TryStar.body);
         VISIT_SEQ(st, excepthandler, s->v.TryStar.handlers);
         VISIT_SEQ(st, stmt, s->v.TryStar.orelse);
         VISIT_SEQ(st, stmt, s->v.TryStar.finalbody);
+        LEAVE_TRY_BLOCK(st);
         LEAVE_CONDITIONAL_BLOCK(st);
         break;
     }
@@ -2098,9 +2139,29 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
             VISIT(st, expr, s->v.Assert.msg);
         break;
     case Import_kind:
+        if (s->v.Import.is_lazy) {
+            if (!check_lazy_import_context(st, s, "import")) {
+                return 0;
+            }
+        }
         VISIT_SEQ(st, alias, s->v.Import.names);
         break;
     case ImportFrom_kind:
+        if (s->v.ImportFrom.is_lazy) {
+            if (!check_lazy_import_context(st, s, "from ... import")) {
+                return 0;
+            }
+
+            /* Check for import * */
+            for (Py_ssize_t i = 0; i < asdl_seq_LEN(s->v.ImportFrom.names); i++) {
+                alias_ty alias = (alias_ty)asdl_seq_GET(s->v.ImportFrom.names, i);
+                if (alias->name && _PyUnicode_EqualToASCIIString(alias->name, "*")) {
+                    PyErr_SetString(PyExc_SyntaxError, "lazy from ... import * is not allowed");
+                    SET_ERROR_LOCATION(st->st_filename, LOCATION(s));
+                    return 0;
+                }
+            }
+        }
         VISIT_SEQ(st, alias, s->v.ImportFrom.names);
         if (!check_import_from(st, s)) {
             return 0;