Add lazy import filter

Include/import.h

@@ -88,6 +88,16 @@ PyAPI_FUNC(int) PyImport_AppendInittab(
     PyObject* (*initfunc)(void)
     );
 
+typedef enum  {
+    PyLazyImportsMode_Default,
+    PyLazyImportsMode_ForcedOff,
+    PyLazyImportsMode_ForcedOn,
+} PyImport_LazyImportsMode;
+
+PyAPI_FUNC(int) PyImport_SetLazyImports(PyImport_LazyImportsMode mode, PyObject *filter);
+
+PyAPI_FUNC(PyImport_LazyImportsMode) PyImport_LazyImportsEnabled(void);
+
 #ifndef Py_LIMITED_API
 #  define Py_CPYTHON_IMPORT_H
 #  include "cpython/import.h"

Include/internal/pycore_ceval.h

@@ -298,7 +298,7 @@ PyAPI_FUNC(void) _PyEval_FormatExcUnbound(PyThreadState *tstate, PyCodeObject *c
 PyAPI_FUNC(void) _PyEval_FormatKwargsError(PyThreadState *tstate, PyObject *func, PyObject *kwargs);
 PyAPI_FUNC(PyObject *) _PyEval_ImportFrom(PyThreadState *, PyObject *, PyObject *);
 PyAPI_FUNC(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);
 PyAPI_FUNC(PyObject *) _PyEval_LazyImportFrom(PyThreadState *tstate, PyObject *v, PyObject *name);
 PyAPI_FUNC(PyObject *) _PyEval_ImportName(PyThreadState *tstate, PyObject *builtins, PyObject *globals, PyObject *locals,
             PyObject *name, PyObject *fromlist, PyObject *level);

Include/internal/pycore_global_objects_fini_generated.h

@@ -1851,6 +1851,7 @@ _PyStaticObjects_CheckRefcnt(PyInterpreterState *interp) {
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(last_type));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(last_value));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(latin1));
+    _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(lazy));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(leaf_size));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(legacy));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(len));

Include/internal/pycore_global_strings.h

@@ -574,6 +574,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(last_type)
         STRUCT_FOR_ID(last_value)
         STRUCT_FOR_ID(latin1)
+        STRUCT_FOR_ID(lazy)
         STRUCT_FOR_ID(leaf_size)
         STRUCT_FOR_ID(legacy)
         STRUCT_FOR_ID(len)

Include/internal/pycore_interp_structs.h

@@ -314,6 +314,8 @@ struct _import_state {
 #endif
     PyObject *import_func;
     PyObject *lazy_import_func;
+    int lazy_imports_mode;
+    PyObject *lazy_imports_filter;
     /* The global import lock. */
     _PyRecursiveMutex lock;
     /* diagnostic info in PyImport_ImportModuleLevelObject() */

Include/internal/pycore_runtime_init_generated.h

@@ -1849,6 +1849,7 @@ extern "C" {
     INIT_ID(last_type), \
     INIT_ID(last_value), \
     INIT_ID(latin1), \
+    INIT_ID(lazy), \
     INIT_ID(leaf_size), \
     INIT_ID(legacy), \
     INIT_ID(len), \

Include/internal/pycore_unicodeobject_generated.h

@@ -2084,6 +2084,10 @@ _PyUnicode_InitStaticStrings(PyInterpreterState *interp) {
     _PyUnicode_InternStatic(interp, &string);
     assert(_PyUnicode_CheckConsistency(string, 1));
     assert(PyUnicode_GET_LENGTH(string) != 1);
+    string = &_Py_ID(lazy);
+    _PyUnicode_InternStatic(interp, &string);
+    assert(_PyUnicode_CheckConsistency(string, 1));
+    assert(PyUnicode_GET_LENGTH(string) != 1);
     string = &_Py_ID(leaf_size);
     _PyUnicode_InternStatic(interp, &string);
     assert(_PyUnicode_CheckConsistency(string, 1));

Lib/importlib/__init__.py

@@ -59,7 +59,7 @@
 
 from ._bootstrap import __import__
 
-from _imp import lazy_import
+from _imp import lazy_import, set_lazy_imports
 
 
 def invalidate_caches():

Lib/test/test_import/__init__.py

@@ -2548,6 +2548,8 @@ def tearDown(self):
             if key.startswith('test.test_import.data.lazy_imports'):
                 del sys.modules[key]
 
+        importlib.set_lazy_imports(None, None)
+
     def test_basic_unused(self):
         try:
             import test.test_import.data.lazy_imports.basic_unused
@@ -2564,6 +2566,54 @@ def test_basic_used(self):
 
         self.assertTrue("test.test_import.data.lazy_imports.basic2" in sys.modules)
 
+    def test_global_off(self):
+        try:
+            import test.test_import.data.lazy_imports.global_off
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertTrue("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
+    def test_global_on(self):
+        try:
+            import test.test_import.data.lazy_imports.global_on
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertFalse("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
+    def test_global_filter(self):
+        try:
+            import test.test_import.data.lazy_imports.global_filter
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertTrue("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
+    def test_global_filter_true(self):
+        try:
+            import test.test_import.data.lazy_imports.global_filter_true
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertFalse("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
+    def test_global_filter_from(self):
+        try:
+            import test.test_import.data.lazy_imports.global_filter
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertTrue("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
+    def test_global_filter_from_true(self):
+        try:
+            import test.test_import.data.lazy_imports.global_filter_true
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertFalse("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
 
 class TestSinglePhaseSnapshot(ModuleSnapshot):
     """A representation of a single-phase init module for testing.

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

@@ -0,0 +1,10 @@
+import importlib
+
+def filter(module_name, imported_name, from_list):
+    assert module_name == __name__
+    assert imported_name == "test.test_import.data.lazy_imports.basic2"
+    return False
+
+importlib.set_lazy_imports(None, filter)
+
+lazy import test.test_import.data.lazy_imports.basic2 as basic2

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

@@ -0,0 +1,11 @@
+import importlib
+
+def filter(module_name, imported_name, from_list):
+    assert module_name == __name__
+    assert imported_name == "test.test_import.data.lazy_imports.basic2"
+    assert from_list == ['f']
+    return False
+
+importlib.set_lazy_imports(None, filter)
+
+lazy from import test.test_import.data.lazy_imports.basic2 import f

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

@@ -0,0 +1,11 @@
+import importlib
+
+def filter(module_name, imported_name, from_list):
+    assert module_name == __name__
+    assert imported_name == "test.test_import.data.lazy_imports.basic2"
+    assert from_list == ['f']
+    return True
+
+importlib.set_lazy_imports(None, filter)
+
+lazy from import test.test_import.data.lazy_imports.basic2 import f

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

@@ -0,0 +1,10 @@
+import importlib
+
+def filter(module_name, imported_name, from_list):
+    assert module_name == __name__
+    assert imported_name == "test.test_import.data.lazy_imports.basic2"
+    return True
+
+importlib.set_lazy_imports(None, filter)
+
+lazy import test.test_import.data.lazy_imports.basic2 as basic2

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

@@ -0,0 +1,5 @@
+import importlib
+
+importlib.set_lazy_imports(False)
+
+lazy import test.test_import.data.lazy_imports.basic2 as basic2

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

@@ -0,0 +1,5 @@
+import importlib
+
+importlib.set_lazy_imports(True)
+
+import test.test_import.data.lazy_imports.basic2 as basic2

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 */
 };