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: D57794109 fbshipit-source-id: 9e5c3ecd31853929f4c5bdc9bf3994d38d8d9c36
facebookincubator · Jun 12, 2024 · 96342a1 · 96342a1
1 parent 4dacb0b
commit 96342a1
Show file tree

Hide file tree

Showing 11 changed files with 151 additions and 31 deletions.
diff --git a/Include/cpython/import.h b/Include/cpython/import.h
@@ -31,7 +31,7 @@ PyAPI_FUNC(int) _PyImport_IsLazyImportsActive(PyThreadState *tstate);
 
 PyAPI_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);
 

diff --git a/Include/internal/pycore_global_objects_fini_generated.h b/Include/internal/pycore_global_objects_fini_generated.h
diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h
@@ -389,6 +389,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(dst_dir_fd)
         STRUCT_FOR_ID(duration)
         STRUCT_FOR_ID(e)
+        STRUCT_FOR_ID(eager)
         STRUCT_FOR_ID(eager_start)
         STRUCT_FOR_ID(effective_ids)
         STRUCT_FOR_ID(element_factory)

diff --git a/Include/internal/pycore_interp.h b/Include/internal/pycore_interp.h
@@ -234,6 +234,7 @@ struct _is {
     // Lazy Imports
     int lazy_imports;  /* whether lazy imports was enabled at runtime */
     PyObject *lazy_import_verbose_seen;
+    PyObject *excluding_modules;
     PyObject *eager_imports;
     PyObject *lazy_modules;
 };

diff --git a/Include/internal/pycore_runtime_init_generated.h b/Include/internal/pycore_runtime_init_generated.h
diff --git a/Include/internal/pycore_unicodeobject_generated.h b/Include/internal/pycore_unicodeobject_generated.h
diff --git a/Lib/importlib/__init__.py b/Lib/importlib/__init__.py
@@ -74,13 +74,17 @@ def 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.
     """
-    return _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):

diff --git a/Lib/test/lazyimports/set_lazy_imports_eager.py b/Lib/test/lazyimports/set_lazy_imports_eager.py
@@ -0,0 +1,38 @@
+# Copyright (c) Meta, Inc. and its affiliates. All Rights Reserved
+# File added for Lazy Imports
+
+import self
+if self._lazy_imports:
+    self.skipTest("Test relevant only when running with global lazy imports disabled")
+
+import importlib
+
+importlib.set_lazy_imports(eager=[
+    "test.lazyimports.data.metasyntactic.foo",
+    "test.lazyimports.data.metasyntactic.waldo",
+    "test.lazyimports.data.metasyntactic.plugh.Plugh",
+])
+
+import test.lazyimports.data.metasyntactic.foo as foo
+self.assertFalse(importlib.is_lazy_import(globals(), "foo"))  # should be eager
+
+from test.lazyimports.data.metasyntactic.foo import bar
+self.assertFalse(importlib.is_lazy_import(globals(), "bar"))  # maybe this should have been lazy?
+
+from test.lazyimports.data.metasyntactic.waldo import Waldo
+self.assertFalse(importlib.is_lazy_import(globals(), "Waldo"))  # maybe this should have been lazy?
+
+import test.lazyimports.data.metasyntactic.waldo.fred as fred
+self.assertTrue(importlib.is_lazy_import(globals(), "fred"))  # this should be lazy
+
+from test.lazyimports.data.metasyntactic.waldo.fred import Fred
+self.assertTrue(importlib.is_lazy_import(globals(), "Fred"))  # this should be lazy
+
+from test.lazyimports.data.metasyntactic.waldo import fred
+self.assertFalse(importlib.is_lazy_import(globals(), "fred"))  # maybe this should have been lazy?
+
+import test.lazyimports.data.metasyntactic.plugh as plugh
+self.assertTrue(importlib.is_lazy_import(globals(), "plugh"))  # this should be lazy
+
+from test.lazyimports.data.metasyntactic.plugh import Plugh
+self.assertFalse(importlib.is_lazy_import(globals(), "Plugh"))  # explicitly eager
diff --git a/Python/clinic/import.c.h b/Python/clinic/import.c.h
diff --git a/Python/import.c b/Python/import.c
@@ -2873,7 +2873,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);
@@ -2883,7 +2883,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;
@@ -2898,7 +2932,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 (true) {
         Py_ssize_t dot = PyUnicode_FindChar(name, '.', 0, PyUnicode_GET_LENGTH(name), -1);
         if (dot < 0) {
@@ -2913,9 +2947,15 @@ add_lazy_modules(PyThreadState *tstate, PyObject *builtins, PyObject *name)
         if (child == 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 excluded, load eagerly */
+                goto end;
+            }
         }
         filter = NULL;
         lazy_submodules = PyDict_GetItemWithError(lazy_modules, parent);
@@ -3052,7 +3092,7 @@ _PyImport_LazyImportName(PyThreadState *tstate, PyObject *builtins, PyObject *gl
         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;
     }
@@ -4382,16 +4422,17 @@ _imp_source_hash_impl(PyObject *module, long key, Py_buffer *source)
 }
 
 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) {
@@ -4405,20 +4446,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);
         }
     }
 
@@ -4496,6 +4556,7 @@ _imp._set_lazy_imports
     enabled: object = True
     /
     excluding: object = NULL
+    eager: object = NULL
 
 Programmatic API for enabling lazy imports at runtime.
 
@@ -4505,10 +4566,10 @@ within modules whose full name is present in the container will be eager.
 
 static PyObject *
 _imp__set_lazy_imports_impl(PyObject *module, PyObject *enabled,
-                            PyObject *excluding)
-/*[clinic end generated code: output=bb6e4196f8cf4569 input=bfae9176e3b98393]*/
+                            PyObject *excluding, PyObject *eager)
+/*[clinic end generated code: output=77575489d11f5806 input=45d71c8b9ad23f16]*/
 {
-    return PyImport_SetLazyImports(enabled, excluding);
+    return PyImport_SetLazyImports(enabled, excluding, eager);
 }
 
 /*[clinic input]
@@ -4542,9 +4603,9 @@ is_lazy_imports_active(PyInterpreterState *interp, _PyInterpreterFrame *frame)
             } else {
                 PyObject *modname  = PyDict_GetItem(frame->f_globals, &_Py_ID(__name__));
                 if (modname != NULL && modname != Py_None) {
-                    PyObject *filter = interp->eager_imports;
+                    PyObject *filter = interp->excluding_modules;
                     if (filter != NULL && PySequence_Contains(filter, modname)) {
-                        lazy_imports = 0;  /* Check imports explicitely set as eager */
+                        lazy_imports = 0;  /* Module explicitly excluded from lazy importing */
                     }
                     PyDict_SetItem(frame->f_globals, &_Py_ID(__lazy_imports_enabled__), lazy_imports ? Py_True : Py_False);
                 } else {

diff --git a/Python/pystate.c b/Python/pystate.c
@@ -710,6 +710,7 @@ init_interpreter(PyInterpreterState *interp,
     interp->f_opcode_trace_set = false;
     interp->lazy_imports = -1;
     interp->lazy_import_verbose_seen = NULL;
+    interp->excluding_modules = NULL;
     interp->eager_imports = NULL;
     interp->lazy_modules = NULL;
 
@@ -881,6 +882,7 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate)
     Py_CLEAR(interp->sysdict_copy);
     Py_CLEAR(interp->builtins_copy);
     Py_CLEAR(interp->lazy_import_verbose_seen);
+    Py_CLEAR(interp->excluding_modules);
     Py_CLEAR(interp->eager_imports);
     Py_CLEAR(interp->lazy_modules);
     Py_CLEAR(interp->dict);