Add __lazy_import__, check sys.modules before import

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 *name, PyObject *fromlist, PyObject *level);
+            PyObject *locals, PyObject *name, PyObject *fromlist, PyObject *level);
 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

@@ -1432,6 +1432,7 @@ _PyStaticObjects_CheckRefcnt(PyInterpreterState *interp) {
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__iter__));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__itruediv__));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__ixor__));
+    _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__lazy_import__));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__le__));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__len__));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__length_hint__));

Include/internal/pycore_global_strings.h

@@ -155,6 +155,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(__iter__)
         STRUCT_FOR_ID(__itruediv__)
         STRUCT_FOR_ID(__ixor__)
+        STRUCT_FOR_ID(__lazy_import__)
         STRUCT_FOR_ID(__le__)
         STRUCT_FOR_ID(__len__)
         STRUCT_FOR_ID(__length_hint__)

Include/internal/pycore_import.h

@@ -35,6 +35,10 @@ extern PyObject *
 _PyImport_ResolveName(PyThreadState *tstate, PyObject *name, PyObject *globals, int level);
 extern PyObject *
 _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import);
+extern PyObject *
+_PyImport_LazyImportModuleLevelObject(PyThreadState *tstate, PyObject *name, PyObject *builtins, PyObject *globals,
+                                      PyObject *locals, PyObject *fromlist,
+                                      int level);
 
 
 #ifdef HAVE_DLOPEN
@@ -79,6 +83,10 @@ extern int _PyImport_IsDefaultImportFunc(
         PyInterpreterState *interp,
         PyObject *func);
 
+extern int _PyImport_IsDefaultLazyImportFunc(
+        PyInterpreterState *interp,
+        PyObject *func);
+
 extern PyObject * _PyImport_GetImportlibLoader(
         PyInterpreterState *interp,
         const char *loader_name);

Include/internal/pycore_interp_structs.h

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

Include/internal/pycore_runtime_init_generated.h

@@ -1430,6 +1430,7 @@ extern "C" {
     INIT_ID(__iter__), \
     INIT_ID(__itruediv__), \
     INIT_ID(__ixor__), \
+    INIT_ID(__lazy_import__), \
     INIT_ID(__le__), \
     INIT_ID(__len__), \
     INIT_ID(__length_hint__), \

Include/internal/pycore_unicodeobject_generated.h

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

Lib/test/test_import/__init__.py

@@ -2540,6 +2540,31 @@ def test_disallowed_reimport(self):
         self.assertIsNot(excsnap, None)
 
 
+class LazyImportTests(unittest.TestCase):
+    def tearDown(self):
+        """Make sure no modules pre-exist in sys.modules which are being used to
+        test."""
+        for key in list(sys.modules.keys()):
+            if key.startswith('test.test_import.data.lazy_imports'):
+                del sys.modules[key]
+
+    def test_basic_unused(self):
+        try:
+            import test.test_import.data.lazy_imports.basic_unused
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertFalse("test.test_import.data.lazy_imports.basic2" in sys.modules)
+
+    def test_basic_used(self):
+        try:
+            import test.test_import.data.lazy_imports.basic_used
+        except ImportError as e:
+            self.fail('lazy import failed')
+
+        self.assertTrue("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/basic2.py

@@ -0,0 +1,4 @@
+def f():
+   pass
+
+x = 42

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

@@ -0,0 +1 @@
+lazy import test.test_import.data.lazy_imports.basic2

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

@@ -0,0 +1,3 @@
+lazy import test.test_import.data.lazy_imports.basic2 as basic2
+
+basic2.f()

Python/bltinmodule.c

@@ -8,6 +8,7 @@
 #include "pycore_fileutils.h"     // _PyFile_Flush
 #include "pycore_floatobject.h"   // _PyFloat_ExactDealloc()
 #include "pycore_interp.h"        // _PyInterpreterState_GetConfig()
+#include "pycore_import.h"        // _PyImport_LazyImportModuleLevelObject  ()
 #include "pycore_long.h"          // _PyLong_CompactValue
 #include "pycore_modsupport.h"    // _PyArg_NoKwnames()
 #include "pycore_object.h"        // _Py_AddToAllObjects()
@@ -287,6 +288,47 @@ builtin___import___impl(PyObject *module, PyObject *name, PyObject *globals,
 }
 
 
+/*[clinic input]
+__lazy_import__ as builtin___lazy_import__
+
+    name: object
+    globals: object(c_default="NULL") = None
+    locals: object(c_default="NULL") = None
+    fromlist: object(c_default="NULL") = ()
+    level: int = 0
+
+Lazily imports a module. 
+
+Returns either the module to be imported or a imp.lazy_module object which
+indicates the module to be lazily imported.
+[clinic start generated code]*/
+
+static PyObject *
+builtin___lazy_import___impl(PyObject *module, PyObject *name,
+                             PyObject *globals, PyObject *locals,
+                             PyObject *fromlist, int level)
+/*[clinic end generated code: output=300f1771094b9e8c input=57123e246d6c36ee]*/
+{
+    PyObject *builtins;
+    PyThreadState *tstate = PyThreadState_GET();
+    if (globals == NULL) {
+        globals = PyEval_GetGlobals();
+    }
+    if (locals == NULL) {
+        locals = globals;
+    }
+
+    if (PyMapping_GetOptionalItem(globals, &_Py_ID(__builtins__), &builtins) < 0) {
+        PyErr_SetString(PyExc_ValueError, "unable to get builtins for lazy import");
+        return NULL;
+    }
+
+    PyObject *res = _PyImport_LazyImportModuleLevelObject(tstate, name, builtins,
+                                                 globals, locals, fromlist, level);
+    Py_DECREF(builtins);
+    return res;
+}
+
 /*[clinic input]
 abs as builtin_abs
 
@@ -3344,6 +3386,7 @@ static PyMethodDef builtin_methods[] = {
     {"__build_class__", _PyCFunction_CAST(builtin___build_class__),
      METH_FASTCALL | METH_KEYWORDS, build_class_doc},
     BUILTIN___IMPORT___METHODDEF
+    BUILTIN___LAZY_IMPORT___METHODDEF
     BUILTIN_ABS_METHODDEF
     BUILTIN_ALL_METHODDEF
     BUILTIN_ANY_METHODDEF

Python/bytecodes.c

@@ -2948,7 +2948,7 @@ dummy_func(
             PyObject *name = GETITEM(FRAME_CO_NAMES, oparg >> 2);
             PyObject *res_o;
             if (oparg & 0x01) {
-                res_o = _PyEval_LazyImportName(tstate, BUILTINS(), GLOBALS(),  name, 
+                res_o = _PyEval_LazyImportName(tstate, BUILTINS(), GLOBALS(), LOCALS(), name, 
                                 PyStackRef_AsPyObjectBorrow(fromlist),
                                 PyStackRef_AsPyObjectBorrow(level));
 

Python/ceval.c

@@ -3027,32 +3027,37 @@ _PyEval_ImportName(PyThreadState *tstate, PyObject *builtins, PyObject *globals,
 
 PyObject *
 _PyEval_LazyImportName(PyThreadState *tstate, PyObject *builtins, PyObject *globals,
-            PyObject *name, PyObject *fromlist, PyObject *level)
+            PyObject *locals, PyObject *name, PyObject *fromlist, PyObject *level)
 {
-    PyObject *res = NULL;
-    PyObject *abs_name = NULL;
-    int ilevel = PyLong_AsInt(level);
-    if (ilevel == -1 && PyErr_Occurred()) {
-        goto error;
+    PyObject *import_func;
+    if (PyMapping_GetOptionalItem(builtins, &_Py_ID(__lazy_import__), &import_func) < 0) {
+        return NULL;
     }
-    if (ilevel > 0) {
-        abs_name = _PyImport_ResolveName(tstate, name, globals, ilevel);
-        if (abs_name == NULL) {
-            goto error;
-        }
-    } else {  /* ilevel == 0 */
-        if (PyUnicode_GET_LENGTH(name) == 0) {
-            PyErr_SetString(PyExc_ValueError, "Empty module name");
-            goto error;
-        }
-        abs_name = name;
-        Py_INCREF(abs_name);
+    if (import_func == NULL) {
+        _PyErr_SetString(tstate, PyExc_ImportError, "__lazy_import__ not found");
+        return NULL;
     }
 
-    // TODO: check sys.modules for module
-    res = _PyLazyImport_New(builtins, abs_name, fromlist);
-error:
-    Py_XDECREF(abs_name);
+    if (locals == NULL) {
+        locals = Py_None;
+    }
+
+    if (_PyImport_IsDefaultLazyImportFunc(tstate->interp, import_func)) {
+        Py_DECREF(import_func);
+
+        int ilevel = PyLong_AsInt(level);
+        if (ilevel == -1 && PyErr_Occurred()) {
+            return NULL;
+        }
+
+        return _PyImport_LazyImportModuleLevelObject(
+            tstate, name, builtins, globals, locals, fromlist, ilevel
+        );
+    }
+
+    PyObject* args[5] = {name, globals, locals, fromlist, level};
+    PyObject *res = PyObject_Vectorcall(import_func, args, 5, NULL);
+    Py_DECREF(import_func);
     return res;
 }
 

Python/clinic/bltinmodule.c.h

@@ -113,6 +113,101 @@ exit:
     return return_value;
 }
 
+PyDoc_STRVAR(builtin___lazy_import____doc__,
+"__lazy_import__($module, /, name, globals=None, locals=None,\n"
+"                fromlist=(), level=0)\n"
+"--\n"
+"\n"
+"Lazily imports a module.\n"
+"\n"
+"Returns either the module to be imported or a imp.lazy_module object which\n"
+"indicates the module to be lazily imported.");
+
+#define BUILTIN___LAZY_IMPORT___METHODDEF    \
+    {"__lazy_import__", _PyCFunction_CAST(builtin___lazy_import__), METH_FASTCALL|METH_KEYWORDS, builtin___lazy_import____doc__},
+
+static PyObject *
+builtin___lazy_import___impl(PyObject *module, PyObject *name,
+                             PyObject *globals, PyObject *locals,
+                             PyObject *fromlist, int level);
+
+static PyObject *
+builtin___lazy_import__(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 5
+    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(name), &_Py_ID(globals), &_Py_ID(locals), &_Py_ID(fromlist), &_Py_ID(level), },
+    };
+    #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[] = {"name", "globals", "locals", "fromlist", "level", NULL};
+    static _PyArg_Parser _parser = {
+        .keywords = _keywords,
+        .fname = "__lazy_import__",
+        .kwtuple = KWTUPLE,
+    };
+    #undef KWTUPLE
+    PyObject *argsbuf[5];
+    Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 1;
+    PyObject *name;
+    PyObject *globals = NULL;
+    PyObject *locals = NULL;
+    PyObject *fromlist = NULL;
+    int level = 0;
+
+    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser,
+            /*minpos*/ 1, /*maxpos*/ 5, /*minkw*/ 0, /*varpos*/ 0, argsbuf);
+    if (!args) {
+        goto exit;
+    }
+    name = args[0];
+    if (!noptargs) {
+        goto skip_optional_pos;
+    }
+    if (args[1]) {
+        globals = args[1];
+        if (!--noptargs) {
+            goto skip_optional_pos;
+        }
+    }
+    if (args[2]) {
+        locals = args[2];
+        if (!--noptargs) {
+            goto skip_optional_pos;
+        }
+    }
+    if (args[3]) {
+        fromlist = args[3];
+        if (!--noptargs) {
+            goto skip_optional_pos;
+        }
+    }
+    level = PyLong_AsInt(args[4]);
+    if (level == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+skip_optional_pos:
+    return_value = builtin___lazy_import___impl(module, name, globals, locals, fromlist, level);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(builtin_abs__doc__,
 "abs($module, number, /)\n"
 "--\n"
@@ -1274,4 +1369,4 @@ builtin_issubclass(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=c0b72519622c849e input=a9049054013a1b77]*/
+/*[clinic end generated code: output=65b2a6dfb50b64bc input=a9049054013a1b77]*/

Python/executor_cases.c.h

@@ -4149,7 +4149,7 @@
             PyObject *res_o;
             if (oparg & 0x01) {
                 _PyFrame_SetStackPointer(frame, stack_pointer);
-                res_o = _PyEval_LazyImportName(tstate, BUILTINS(), GLOBALS(),  name,
+                res_o = _PyEval_LazyImportName(tstate, BUILTINS(), GLOBALS(), LOCALS(), name,
                     PyStackRef_AsPyObjectBorrow(fromlist),
                     PyStackRef_AsPyObjectBorrow(level));
                 stack_pointer = _PyFrame_GetStackPointer(frame);

Python/generated_cases.c.h

@@ -6225,7 +6225,7 @@
             PyObject *res_o;
             if (oparg & 0x01) {
                 _PyFrame_SetStackPointer(frame, stack_pointer);
-                res_o = _PyEval_LazyImportName(tstate, BUILTINS(), GLOBALS(),  name,
+                res_o = _PyEval_LazyImportName(tstate, BUILTINS(), GLOBALS(), LOCALS(), name,
                     PyStackRef_AsPyObjectBorrow(fromlist),
                     PyStackRef_AsPyObjectBorrow(level));
                 stack_pointer = _PyFrame_GetStackPointer(frame);

Python/import.c

@@ -99,6 +99,9 @@ static struct _inittab *inittab_copy = NULL;
 #define IMPORT_FUNC(interp) \
     (interp)->imports.import_func
 
+#define LAZY_IMPORT_FUNC(interp) \
+    (interp)->imports.lazy_import_func
+
 #define IMPORT_LOCK(interp) \
     (interp)->imports.lock
 
@@ -3400,6 +3403,12 @@ _PyImport_InitDefaultImportFunc(PyInterpreterState *interp)
         return -1;
     }
     IMPORT_FUNC(interp) = import_func;
+
+    // Get the __lazy_import__ function
+    if (PyDict_GetItemStringRef(interp->builtins, "__lazy_import__", &import_func) <= 0) {
+        return -1;
+    }
+    LAZY_IMPORT_FUNC(interp) = import_func;
     return 0;
 }
 
@@ -3409,6 +3418,11 @@ _PyImport_IsDefaultImportFunc(PyInterpreterState *interp, PyObject *func)
     return func == IMPORT_FUNC(interp);
 }
 
+int
+_PyImport_IsDefaultLazyImportFunc(PyInterpreterState *interp, PyObject *func)
+{
+    return func == LAZY_IMPORT_FUNC(interp);
+}
 
 /* Import a module, either built-in, frozen, or external, and return
    its module object WITH INCREMENTED REFERENCE COUNT */
@@ -3685,8 +3699,6 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
     PyObject *fromlist = NULL;
     assert(lazy_import != NULL);
     assert(PyLazyImport_CheckExact(lazy_import));
-    PyObject *state_dict = _PyThreadState_GetDict(tstate);
-    assert(state_dict != NULL);
 
     PyLazyImportObject *lz = (PyLazyImportObject *)lazy_import;
 
@@ -3821,6 +3833,19 @@ import_find_and_load(PyThreadState *tstate, PyObject *abs_name)
 #undef accumulated
 }
 
+PyObject *
+get_abs_name(PyThreadState *tstate, PyObject *name, PyObject *globals, int level)
+{
+    if (level > 0) {
+        return resolve_name(tstate, name, globals, level);
+    }
+    if (PyUnicode_GET_LENGTH(name) == 0) {
+        _PyErr_SetString(tstate, PyExc_ValueError, "Empty module name");
+        return NULL;
+    }
+    return Py_NewRef(name);
+}
+
 PyObject *
 PyImport_ImportModuleLevelObject(PyObject *name, PyObject *globals,
                                  PyObject *locals, PyObject *fromlist,
@@ -3852,17 +3877,9 @@ PyImport_ImportModuleLevelObject(PyObject *name, PyObject *globals,
         goto error;
     }
 
-    if (level > 0) {
-        abs_name = resolve_name(tstate, name, globals, level);
-        if (abs_name == NULL)
-            goto error;
-    }
-    else {  /* level == 0 */
-        if (PyUnicode_GET_LENGTH(name) == 0) {
-            _PyErr_SetString(tstate, PyExc_ValueError, "Empty module name");
-            goto error;
-        }
-        abs_name = Py_NewRef(name);
+    abs_name = get_abs_name(tstate, name, globals, level);
+    if (abs_name == NULL) {
+        goto error;
     }
 
     mod = import_get_module(tstate, abs_name);
@@ -3966,6 +3983,28 @@ PyImport_ImportModuleLevelObject(PyObject *name, PyObject *globals,
     return final_mod;
 }
 
+PyObject *
+_PyImport_LazyImportModuleLevelObject(PyThreadState *tstate,
+                                      PyObject *name, PyObject *builtins,
+                                      PyObject *globals, PyObject *locals,
+                                      PyObject *fromlist, int level)
+{
+    PyObject *abs_name = get_abs_name(tstate, name, globals, level);
+    if (abs_name == NULL) {
+        return NULL;
+    }
+
+    PyObject *mod = PyImport_GetModule(abs_name);
+    if (mod != NULL) {
+        Py_DECREF(abs_name);
+        return mod;
+    }
+
+    PyObject *res = _PyLazyImport_New(builtins, abs_name, fromlist);
+    Py_DECREF(abs_name);
+    return res;
+}
+
 PyObject *
 PyImport_ImportModuleLevel(const char *name, PyObject *globals, PyObject *locals,
                            PyObject *fromlist, int level)
@@ -4172,6 +4211,7 @@ _PyImport_ClearCore(PyInterpreterState *interp)
     Py_CLEAR(MODULES_BY_INDEX(interp));
     Py_CLEAR(IMPORTLIB(interp));
     Py_CLEAR(IMPORT_FUNC(interp));
+    Py_CLEAR(LAZY_IMPORT_FUNC(interp));
 }
 
 void