Add support for eager imports

Summary:
We had support for having modules excluded from lazy importing things. This was set via `importlib.set_lazy_imports(excluding=...)`.

This diff adds support for pinning eager imports explicitly via `importlib.set_lazy_imports(eager=...)`.

The difference between the two is that with the former, any imports inside modules in the `excluding` container would be eagerly imported. With the later, anything listed in the `eager` container would always be eagerly imported, no
matter where it's imported. Modules in this container have to be exact matches to whatever is being imported.

One caveat to be aware is that `importlib.set_lazy_imports(eager={"foo.bar.baz"})` will make `import foo.bar.baz` or `from foo.bar import baz` eager, but not `import foo.bar.baz.plugh`, even if that import would imply importing `foo.bar.baz`.

Differential Revision: D57791527

fbshipit-source-id: 56e80458c73f9ad4808e3eb6368f4ec061a57479

Include/import.h

@@ -59,7 +59,7 @@ PyAPI_FUNC(PyObject *) PyImport_ImportModuleLevel(
 #if !defined(Py_LIMITED_API) || Py_LIMITED_API+0 >= 0x03050000
 CiAPI_FUNC(int) PyImport_IsLazyImportsEnabled(void);
 PyAPI_FUNC(PyObject *) PyImport_SetLazyImports(
-    PyObject *enabled, PyObject *excluding);
+    PyObject *enabled, PyObject *excluding, PyObject *eager);
 PyAPI_FUNC(PyObject *) PyImport_SetLazyImportsInModule(
     PyObject *enabled);
 PyObject *

Include/internal/pycore_interp.h

@@ -251,6 +251,7 @@ struct _is {
     PyObject *builtins;
     // importlib module
     PyObject *importlib;
+    PyObject *excluding_modules;
     PyObject *eager_imports;
 
     /* Used in Modules/_threadmodule.c. */

Lib/importlib/__init__.py

@@ -67,13 +67,17 @@
 hydrate_lazy_objects = _imp.hydrate_lazy_objects
 
 
-def set_lazy_imports(enable = True, *, excluding = None):
+def set_lazy_imports(enable = True, /, excluding = None, eager = None):
     """Programmatic API for enabling lazy imports at runtime.
 
-    The optional argument `excluding` can be any container of strings; all imports
-    within modules whose full name is present in the container will be eager.
+    The optional argument `excluding` can be any container of strings; all the
+    modules whose full name is present in the container will be excluded from
+    having any lazy imports, so all the imports within such modules will be eager.
+
+    The optional argument `eager` can be any container of strings; all imports for
+    which the import full name is present in the container will be imported eagerly.
     """
-    _imp._set_lazy_imports(enable, excluding=excluding)
+    return _imp._set_lazy_imports(enable, excluding=excluding, eager=eager)
 
 
 def enable_lazy_imports_in_module(enable = True):

Lib/test/lazyimports/immediate_set_lazy_import.py

@@ -9,9 +9,14 @@
 
 from test.lazyimports.data.metasyntactic import foo
 
+importlib.set_lazy_imports(eager=["test.lazyimports.data.metasyntactic.plugh"])
+
+from test.lazyimports.data.metasyntactic import plugh
+
 importlib.set_lazy_imports(excluding=["test.lazyimports.immediate_set_lazy_import"])
 
 from test.lazyimports.data.metasyntactic import waldo
 
 self.assertTrue(importlib.is_lazy_import(globals(), "foo"))
+self.assertFalse(importlib.is_lazy_import(globals(), "plugh"))
 self.assertFalse(importlib.is_lazy_import(globals(), "waldo"))

Python/import.c

@@ -1624,7 +1624,7 @@ add_lazy_submodule(PyObject *module, PyObject *name)
 }
 
 static int
-add_lazy_modules(PyThreadState *tstate, PyObject *builtins, PyObject *name)
+add_lazy_modules(PyThreadState *tstate, PyObject *builtins, PyObject *name, PyObject *fromlist)
 {
     int ret = 1;
     assert(tstate->interp->lazy_modules != NULL);
@@ -1634,7 +1634,41 @@ add_lazy_modules(PyThreadState *tstate, PyObject *builtins, PyObject *name)
     PyObject *child = NULL;
     PyObject *parent_module = NULL;
     PyObject *parent_dict = NULL;
-    PyObject *lazy_submodules = PyDict_GetItemWithError(lazy_modules, name);
+    PyObject *lazy_submodules;
+
+    if (tstate->interp->eager_imports != NULL) {
+        int found = PySequence_Contains(tstate->interp->eager_imports, name);
+        if (found < 0) {
+            goto error;
+        }
+        if (found) {
+            ret = 0; /* If the module is flagged as eager import, load eagerly */
+            goto end;
+        }
+        if (fromlist != NULL && fromlist != Py_None) {
+            assert(PyTuple_CheckExact(fromlist));
+            Py_ssize_t size = PyTuple_GET_SIZE(fromlist);
+            for (Py_ssize_t i = 0; i < size; ++i) {
+                PyObject* item = PyTuple_GET_ITEM(fromlist, i);
+                assert(PyUnicode_Check(item));
+                PyObject *from_name = PyUnicode_FromFormat("%U.%U", name, item);
+                if (from_name == NULL) {
+                    goto error;
+                }
+                found = PySequence_Contains(tstate->interp->eager_imports, from_name);
+                Py_DECREF(from_name);
+                if (found < 0) {
+                    goto error;
+                }
+                if (found) {
+                    ret = 0; /* If the module is flagged as eager import, load eagerly */
+                    goto end;
+                }
+            }
+        }
+    }
+
+    lazy_submodules = PyDict_GetItemWithError(lazy_modules, name);
     if (lazy_submodules == NULL) {
         if (PyErr_Occurred()) {
             goto error;
@@ -1649,7 +1683,7 @@ add_lazy_modules(PyThreadState *tstate, PyObject *builtins, PyObject *name)
         }
         Py_DECREF(lazy_submodules);
     }
-    PyObject *filter = tstate->interp->eager_imports;
+    PyObject *filter = tstate->interp->excluding_modules;
     while (1) {
         Py_ssize_t dot = PyUnicode_FindChar(name, '.', 0, PyUnicode_GET_LENGTH(name), -1);
         if (dot < 0) {
@@ -1659,9 +1693,15 @@ add_lazy_modules(PyThreadState *tstate, PyObject *builtins, PyObject *name)
         if (parent == NULL) {
             goto error;
         }
-        if (filter != NULL && PySequence_Contains(filter, parent)) {
-            ret = 0; /* If the direct parent is eager, load eagerly */
-            goto end;
+        if (filter != NULL) {
+            int found = PySequence_Contains(filter, parent);
+            if (found < 0) {
+                goto error;
+            }
+            if (found) {
+                ret = 0; /* If the direct parent is eager, load eagerly */
+                goto end;
+            }
         }
         filter = NULL;
         Py_XDECREF(child);
@@ -1808,7 +1848,7 @@ _PyImport_LazyImportName(PyObject *builtins, PyObject *globals, PyObject *locals
         Py_INCREF(abs_name);
     }
 
-    int lazy = add_lazy_modules(tstate, builtins, abs_name);
+    int lazy = add_lazy_modules(tstate, builtins, abs_name, fromlist);
     if (lazy < 0) {
         goto error;
     }
@@ -2657,16 +2697,17 @@ closest_module_frame(PyFrameObject *frame)
 }
 
 PyObject *
-PyImport_SetLazyImports(PyObject *enabled, PyObject *excluding)
+PyImport_SetLazyImports(PyObject *enabled, PyObject *excluding, PyObject *eager)
 {
     PyObject *result = NULL;
     PyInterpreterState *interp = _PyInterpreterState_GET();
     assert(interp != NULL);
     assert(interp->lazy_imports != -1);
 
     result = PyTuple_Pack(
-        2,
+        3,
         interp->lazy_imports ? Py_True : Py_False,
+        interp->excluding_modules == NULL ? Py_None : interp->excluding_modules,
         interp->eager_imports == NULL ? Py_None : interp->eager_imports
     );
     if (result == NULL) {
@@ -2680,20 +2721,39 @@ PyImport_SetLazyImports(PyObject *enabled, PyObject *excluding)
 
     if (excluding != NULL) {
         if (Py_IsNone(excluding)) {
+            Py_XDECREF(interp->excluding_modules);
+            interp->excluding_modules = NULL;
+        } else {
+            PyObject *empty = PyUnicode_New(0, 0);
+            if (empty == NULL) {
+                goto error;
+            }
+            if (PySequence_Contains(excluding, empty) == -1) {
+                Py_DECREF(empty);
+                goto error;
+            }
+            Py_DECREF(empty);
+            Py_XDECREF(interp->excluding_modules);
+            interp->excluding_modules = Py_NewRef(excluding);
+        }
+    }
+
+    if (eager != NULL) {
+        if (Py_IsNone(eager)) {
             Py_XDECREF(interp->eager_imports);
             interp->eager_imports = NULL;
         } else {
             PyObject *empty = PyUnicode_New(0, 0);
             if (empty == NULL) {
                 goto error;
             }
-            if (PySequence_Contains(excluding, empty) == -1) {
+            if (PySequence_Contains(eager, empty) == -1) {
                 Py_DECREF(empty);
                 goto error;
             }
             Py_DECREF(empty);
             Py_XDECREF(interp->eager_imports);
-            interp->eager_imports = Py_NewRef(excluding);
+            interp->eager_imports = Py_NewRef(eager);
         }
     }
 
@@ -2777,6 +2837,7 @@ _imp._set_lazy_imports
     enabled: object = True
     /
     excluding: object = NULL
+    eager: object = NULL
 
 Programmatic API for enabling lazy imports at runtime.
 
@@ -2786,10 +2847,10 @@ within which all imports will remain eager.
 
 static PyObject *
 _imp__set_lazy_imports_impl(PyObject *module, PyObject *enabled,
-                            PyObject *excluding)
-/*[clinic end generated code: output=bb6e4196f8cf4569 input=c16f11904690a50f]*/
+                            PyObject *excluding, PyObject *eager)
+/*[clinic end generated code: output=77575489d11f5806 input=6fbccf14b8258afb]*/
 {
-    return PyImport_SetLazyImports(enabled, excluding);
+    return PyImport_SetLazyImports(enabled, excluding, eager);
 }
 
 /*[clinic input]
@@ -3293,7 +3354,7 @@ _PyImport_IsLazyImportsEnabled(PyThreadState *tstate)
                 if (lazy_imports) {
                     PyObject *modname = _PyDict_GetItemIdWithError(frame->f_globals, &PyId___name__);
                     if (modname != NULL && modname != Py_None) {
-                        PyObject *filter = tstate->interp->eager_imports;
+                        PyObject *filter = tstate->interp->excluding_modules;
                         if (filter != NULL && PySequence_Contains(filter, modname)) {
                             lazy_imports = 0;  /* Check imports explicitely set as eager */
                         }

Python/pystate.c

@@ -275,6 +275,7 @@ PyInterpreterState_New(void)
 
     interp->audit_hooks = NULL;
     interp->lazy_imports = 0;
+    interp->excluding_modules = NULL;
     interp->eager_imports = NULL;
     interp->lazy_modules = NULL;
 
@@ -316,6 +317,7 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate)
     Py_CLEAR(interp->builtins_copy);
     Py_CLEAR(interp->importlib);
     Py_CLEAR(interp->import_func);
+    Py_CLEAR(interp->excluding_modules);
     Py_CLEAR(interp->eager_imports);
     Py_CLEAR(interp->lazy_modules);
     Py_CLEAR(interp->dict);