Add lazy import filter

Python/bytecodes.c

@@ -2947,10 +2947,11 @@ dummy_func(
          inst(IMPORT_NAME, (level, fromlist -- res)) {
             PyObject *name = GETITEM(FRAME_CO_NAMES, oparg >> 2);
             PyObject *res_o;
-            if (oparg & 0x01) {
+            if (!(oparg & 0x02)) {
                 res_o = _PyEval_LazyImportName(tstate, BUILTINS(), GLOBALS(), LOCALS(), name, 
                                 PyStackRef_AsPyObjectBorrow(fromlist),
-                                PyStackRef_AsPyObjectBorrow(level));
+                                PyStackRef_AsPyObjectBorrow(level),
+                                oparg & 0x01);
 
             } else {
                 res_o = _PyEval_ImportName(tstate, BUILTINS(), GLOBALS(), LOCALS(), name,

Python/ceval.c

@@ -3027,8 +3027,25 @@ _PyEval_ImportName(PyThreadState *tstate, PyObject *builtins, PyObject *globals,
 
 PyObject *
 _PyEval_LazyImportName(PyThreadState *tstate, PyObject *builtins, PyObject *globals,
-            PyObject *locals, PyObject *name, PyObject *fromlist, PyObject *level)
+            PyObject *locals, PyObject *name, PyObject *fromlist, PyObject *level, int lazy)
 {
+    // Check if global policy overrides the local syntax
+    switch (PyImport_LazyImportsEnabled()) {
+        case PyLazyImportsMode_ForcedOff:
+            lazy = 0;
+            break;
+        case PyLazyImportsMode_ForcedOn:
+            lazy = 1;
+            break;
+        case PyLazyImportsMode_Default:
+            break;
+    }
+
+    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);
+    }
+
     PyObject *import_func;
     if (PyMapping_GetOptionalItem(builtins, &_Py_ID(__lazy_import__), &import_func) < 0) {
         return NULL;

Python/clinic/import.c.h

@@ -622,6 +622,88 @@ exit:
     return return_value;
 }
 
+PyDoc_STRVAR(_imp_set_lazy_imports__doc__,
+"set_lazy_imports($module, enabled=None, /, filter=<unrepresentable>)\n"
+"--\n"
+"\n"
+"Programmatic API for enabling lazy imports at runtime.\n"
+"\n"
+"enabled can be:\n"
+"    None (lazy imports always respect keyword)\n"
+"    False (forced lazy imports off)\n"
+"    True (forced lazy imports on)\n"
+"\n"
+"filter is an optional callable which further disables lazy imports when they\n"
+"would otherwise be enabled. Returns True if the the import is still enabled\n"
+"or False to disable it. The callable is called with:\n"
+"\n"
+"(importing_module_name, imported_module_name, [fromlist])");
+
+#define _IMP_SET_LAZY_IMPORTS_METHODDEF    \
+    {"set_lazy_imports", _PyCFunction_CAST(_imp_set_lazy_imports), METH_FASTCALL|METH_KEYWORDS, _imp_set_lazy_imports__doc__},
+
+static PyObject *
+_imp_set_lazy_imports_impl(PyObject *module, PyObject *enabled,
+                           PyObject *filter);
+
+static PyObject *
+_imp_set_lazy_imports(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE)
+
+    #define NUM_KEYWORDS 1
+    static struct {
+        PyGC_Head _this_is_not_used;
+        PyObject_VAR_HEAD
+        Py_hash_t ob_hash;
+        PyObject *ob_item[NUM_KEYWORDS];
+    } _kwtuple = {
+        .ob_base = PyVarObject_HEAD_INIT(&PyTuple_Type, NUM_KEYWORDS)
+        .ob_hash = -1,
+        .ob_item = { &_Py_ID(filter), },
+    };
+    #undef NUM_KEYWORDS
+    #define KWTUPLE (&_kwtuple.ob_base.ob_base)
+
+    #else  // !Py_BUILD_CORE
+    #  define KWTUPLE NULL
+    #endif  // !Py_BUILD_CORE
+
+    static const char * const _keywords[] = {"", "filter", NULL};
+    static _PyArg_Parser _parser = {
+        .keywords = _keywords,
+        .fname = "set_lazy_imports",
+        .kwtuple = KWTUPLE,
+    };
+    #undef KWTUPLE
+    PyObject *argsbuf[2];
+    Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 0;
+    PyObject *enabled = Py_None;
+    PyObject *filter = NULL;
+
+    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser,
+            /*minpos*/ 0, /*maxpos*/ 2, /*minkw*/ 0, /*varpos*/ 0, argsbuf);
+    if (!args) {
+        goto exit;
+    }
+    if (nargs < 1) {
+        goto skip_optional_posonly;
+    }
+    noptargs--;
+    enabled = args[0];
+skip_optional_posonly:
+    if (!noptargs) {
+        goto skip_optional_pos;
+    }
+    filter = args[1];
+skip_optional_pos:
+    return_value = _imp_set_lazy_imports_impl(module, enabled, filter);
+
+exit:
+    return return_value;
+}
+
 #ifndef _IMP_CREATE_DYNAMIC_METHODDEF
     #define _IMP_CREATE_DYNAMIC_METHODDEF
 #endif /* !defined(_IMP_CREATE_DYNAMIC_METHODDEF) */
@@ -629,4 +711,4 @@ exit:
 #ifndef _IMP_EXEC_DYNAMIC_METHODDEF
     #define _IMP_EXEC_DYNAMIC_METHODDEF
 #endif /* !defined(_IMP_EXEC_DYNAMIC_METHODDEF) */
-/*[clinic end generated code: output=24f597d6b0f3feed input=a9049054013a1b77]*/
+/*[clinic end generated code: output=2c52abee5d118061 input=a9049054013a1b77]*/

Python/executor_cases.c.h

@@ -4147,11 +4147,12 @@
             level = stack_pointer[-2];
             PyObject *name = GETITEM(FRAME_CO_NAMES, oparg >> 2);
             PyObject *res_o;
-            if (oparg & 0x01) {
+            if (!(oparg & 0x02)) {
                 _PyFrame_SetStackPointer(frame, stack_pointer);
                 res_o = _PyEval_LazyImportName(tstate, BUILTINS(), GLOBALS(), LOCALS(), name,
                     PyStackRef_AsPyObjectBorrow(fromlist),
-                    PyStackRef_AsPyObjectBorrow(level));
+                    PyStackRef_AsPyObjectBorrow(level),
+                    oparg & 0x01);
                 stack_pointer = _PyFrame_GetStackPointer(frame);
             } else {
                 _PyFrame_SetStackPointer(frame, stack_pointer);

Python/generated_cases.c.h

@@ -6223,11 +6223,12 @@
             level = stack_pointer[-2];
             PyObject *name = GETITEM(FRAME_CO_NAMES, oparg >> 2);
             PyObject *res_o;
-            if (oparg & 0x01) {
+            if (!(oparg & 0x02)) {
                 _PyFrame_SetStackPointer(frame, stack_pointer);
                 res_o = _PyEval_LazyImportName(tstate, BUILTINS(), GLOBALS(), LOCALS(), name,
                     PyStackRef_AsPyObjectBorrow(fromlist),
-                    PyStackRef_AsPyObjectBorrow(level));
+                    PyStackRef_AsPyObjectBorrow(level),
+                    oparg & 0x01);
                 stack_pointer = _PyFrame_GetStackPointer(frame);
             } else {
                 _PyFrame_SetStackPointer(frame, stack_pointer);

Python/import.c

@@ -108,6 +108,12 @@ static struct _inittab *inittab_copy = NULL;
 #define FIND_AND_LOAD(interp) \
     (interp)->imports.find_and_load
 
+#define LAZY_IMPORTS_MODE(interp) \
+    (interp)->imports.lazy_imports_mode
+
+#define LAZY_IMPORTS_FILTER(interp) \
+    (interp)->imports.lazy_imports_filter
+
 #define _IMPORT_TIME_HEADER(interp)                                           \
     do {                                                                      \
         if (FIND_AND_LOAD((interp)).header) {                                 \
@@ -3994,12 +4000,44 @@ _PyImport_LazyImportModuleLevelObject(PyThreadState *tstate,
         return NULL;
     }
 
+    PyInterpreterState *interp = tstate->interp;
+
     PyObject *mod = PyImport_GetModule(abs_name);
     if (mod != NULL) {
         Py_DECREF(abs_name);
         return mod;
     }
 
+    // Check if the filter disables the lazy import
+    PyObject *filter = LAZY_IMPORTS_FILTER(interp);
+    if (filter != NULL) {
+        PyObject *modname;
+        if (PyDict_GetItemRef(globals, &_Py_ID(__name__), &modname) < 0) {
+            return NULL;
+        } else if (modname == NULL) {
+            modname = Py_None;
+        }
+        PyObject *args[] = {modname, name, fromlist};
+        PyObject *res = PyObject_Vectorcall(
+            filter,
+            args,
+            3,
+            NULL
+        );
+
+        if (res == NULL) {
+            Py_DECREF(abs_name);
+            return NULL;
+        }
+
+        if (!PyObject_IsTrue(res)) {
+            Py_DECREF(abs_name);
+            return PyImport_ImportModuleLevelObject(
+                name, globals, locals, fromlist, level
+            );
+        }
+    }
+
     PyObject *res = _PyLazyImport_New(builtins, abs_name, fromlist);
     Py_DECREF(abs_name);
     return res;
@@ -4348,6 +4386,31 @@ PyImport_ImportModuleAttrString(const char *modname, const char *attrname)
 }
 
 
+int
+PyImport_SetLazyImports(PyImport_LazyImportsMode mode, PyObject *filter)
+{
+    if (filter == Py_None) {
+        filter = NULL;
+    }
+    if (filter != NULL && !PyCallable_Check(filter)) {
+        PyErr_SetString(PyExc_ValueError, "filter provided but is not callable");
+        return -1;
+    }
+
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    LAZY_IMPORTS_MODE(interp) = mode;
+    Py_XSETREF(LAZY_IMPORTS_FILTER(interp), Py_XNewRef(filter));
+    return 0;
+}
+
+/* Checks if lazy imports is globally enabled or disabled. Return 1 when globally
+ * forced on, 0 when globally forced off, or -1 when */
+PyImport_LazyImportsMode PyImport_LazyImportsEnabled(void)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    return LAZY_IMPORTS_MODE(interp);
+}
+
 /**************/
 /* the module */
 /**************/
@@ -4937,6 +5000,51 @@ _imp_source_hash_impl(PyObject *module, long key, Py_buffer *source)
     return PyBytes_FromStringAndSize(hash.data, sizeof(hash.data));
 }
 
+/*[clinic input]
+_imp.set_lazy_imports
+
+    enabled: object = None
+    /
+    filter: object = NULL
+
+Programmatic API for enabling lazy imports at runtime.
+
+enabled can be:
+    None (lazy imports always respect keyword)
+    False (forced lazy imports off)
+    True (forced lazy imports on)
+
+filter is an optional callable which further disables lazy imports when they
+would otherwise be enabled. Returns True if the the import is still enabled
+or False to disable it. The callable is called with:
+
+(importing_module_name, imported_module_name, [fromlist])
+
+[clinic start generated code]*/
+
+static PyObject *
+_imp_set_lazy_imports_impl(PyObject *module, PyObject *enabled,
+                           PyObject *filter)
+/*[clinic end generated code: output=d8d5a848c041edc5 input=00b2334fae4345a3]*/
+{
+    PyImport_LazyImportsMode mode;
+    if (enabled == Py_None) {
+        mode = PyLazyImportsMode_Default;
+    } else if (enabled == Py_False) {
+        mode = PyLazyImportsMode_ForcedOff;
+    } else if (enabled == Py_True) {
+        mode = PyLazyImportsMode_ForcedOn;
+    } else {
+        PyErr_SetString(PyExc_ValueError, "expected None, True or False for enabled mode");
+        return NULL;
+    }
+
+    if (PyImport_SetLazyImports(mode, filter) < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
 
 PyDoc_STRVAR(doc_imp,
 "(Extremely) low-level import machinery bits as used by importlib.");
@@ -4961,6 +5069,7 @@ static PyMethodDef imp_methods[] = {
     _IMP_EXEC_BUILTIN_METHODDEF
     _IMP__FIX_CO_FILENAME_METHODDEF
     _IMP_SOURCE_HASH_METHODDEF
+    _IMP_SET_LAZY_IMPORTS_METHODDEF
     {NULL, NULL}  /* sentinel */
 };
 
@@ -4980,7 +5089,7 @@ imp_module_exec(PyObject *module)
         return -1;
     }
 
-    if (PyModule_AddObjectRef(module, "lazy_import", 
+    if (PyModule_AddObjectRef(module, "lazy_import",
                               (PyObject *)&PyLazyImport_Type) < 0) {
         return -1;
     }