From 96342a124f4b11bb493380833435912a28549f65 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Germ=C3=A1n=20M=C3=A9ndez=20Bravo?= Date: Wed, 12 Jun 2024 12:20:11 -0700 Subject: [PATCH] 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 --- Include/cpython/import.h | 2 +- .../pycore_global_objects_fini_generated.h | 1 + Include/internal/pycore_global_strings.h | 1 + Include/internal/pycore_interp.h | 1 + .../internal/pycore_runtime_init_generated.h | 1 + .../internal/pycore_unicodeobject_generated.h | 3 + Lib/importlib/__init__.py | 12 ++- .../lazyimports/set_lazy_imports_eager.py | 38 ++++++++ Python/clinic/import.c.h | 28 ++++-- Python/import.c | 93 +++++++++++++++---- Python/pystate.c | 2 + 11 files changed, 151 insertions(+), 31 deletions(-) create mode 100644 Lib/test/lazyimports/set_lazy_imports_eager.py diff --git a/Include/cpython/import.h b/Include/cpython/import.h index baeb2f2e402..44d4147bf4f 100644 --- 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 index a46608bf0ef..a985c6405e5 100644 --- a/Include/internal/pycore_global_objects_fini_generated.h +++ b/Include/internal/pycore_global_objects_fini_generated.h @@ -900,6 +900,7 @@ _PyStaticObjects_CheckRefcnt(PyInterpreterState *interp) { _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(dst_dir_fd)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(duration)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(e)); + _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(eager)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(eager_start)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(effective_ids)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(element_factory)); diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h index 32c76eb92fd..b251fc387b5 100644 --- 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 index 67f44faa03c..b801ef95c29 100644 --- 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 index 5833efa8de3..743f666dcd5 100644 --- a/Include/internal/pycore_runtime_init_generated.h +++ b/Include/internal/pycore_runtime_init_generated.h @@ -895,6 +895,7 @@ extern "C" { INIT_ID(dst_dir_fd), \ INIT_ID(duration), \ INIT_ID(e), \ + INIT_ID(eager), \ INIT_ID(eager_start), \ INIT_ID(effective_ids), \ INIT_ID(element_factory), \ diff --git a/Include/internal/pycore_unicodeobject_generated.h b/Include/internal/pycore_unicodeobject_generated.h index 896b309b09c..bd4a895a63a 100644 --- a/Include/internal/pycore_unicodeobject_generated.h +++ b/Include/internal/pycore_unicodeobject_generated.h @@ -1008,6 +1008,9 @@ _PyUnicode_InitStaticStrings(PyInterpreterState *interp) { string = &_Py_ID(e); assert(_PyUnicode_CheckConsistency(string, 1)); _PyUnicode_InternInPlace(interp, &string); + string = &_Py_ID(eager); + assert(_PyUnicode_CheckConsistency(string, 1)); + _PyUnicode_InternInPlace(interp, &string); string = &_Py_ID(eager_start); assert(_PyUnicode_CheckConsistency(string, 1)); _PyUnicode_InternInPlace(interp, &string); diff --git a/Lib/importlib/__init__.py b/Lib/importlib/__init__.py index e3d18047506..a3afa6fe9d6 100644 --- 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 new file mode 100644 index 00000000000..97db942b3a1 --- /dev/null +++ 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 index eb93a8b19a2..2691dc1f85e 100644 --- a/Python/clinic/import.c.h +++ b/Python/clinic/import.c.h @@ -687,7 +687,8 @@ _imp_is_lazy_import(PyObject *module, PyObject *const *args, Py_ssize_t nargs) } PyDoc_STRVAR(_imp__set_lazy_imports__doc__, -"_set_lazy_imports($module, enabled=True, /, excluding=)\n" +"_set_lazy_imports($module, enabled=True, /,\n" +" excluding=, eager=)\n" "--\n" "\n" "Programmatic API for enabling lazy imports at runtime.\n" @@ -700,7 +701,7 @@ PyDoc_STRVAR(_imp__set_lazy_imports__doc__, static PyObject * _imp__set_lazy_imports_impl(PyObject *module, PyObject *enabled, - PyObject *excluding); + PyObject *excluding, PyObject *eager); static PyObject * _imp__set_lazy_imports(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) @@ -708,14 +709,14 @@ _imp__set_lazy_imports(PyObject *module, PyObject *const *args, Py_ssize_t nargs PyObject *return_value = NULL; #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE) - #define NUM_KEYWORDS 1 + #define NUM_KEYWORDS 2 static struct { PyGC_Head _this_is_not_used; PyObject_VAR_HEAD PyObject *ob_item[NUM_KEYWORDS]; } _kwtuple = { .ob_base = PyVarObject_HEAD_INIT(&PyTuple_Type, NUM_KEYWORDS) - .ob_item = { &_Py_ID(excluding), }, + .ob_item = { &_Py_ID(excluding), &_Py_ID(eager), }, }; #undef NUM_KEYWORDS #define KWTUPLE (&_kwtuple.ob_base.ob_base) @@ -724,19 +725,20 @@ _imp__set_lazy_imports(PyObject *module, PyObject *const *args, Py_ssize_t nargs # define KWTUPLE NULL #endif // !Py_BUILD_CORE - static const char * const _keywords[] = {"", "excluding", NULL}; + static const char * const _keywords[] = {"", "excluding", "eager", NULL}; static _PyArg_Parser _parser = { .keywords = _keywords, .fname = "_set_lazy_imports", .kwtuple = KWTUPLE, }; #undef KWTUPLE - PyObject *argsbuf[2]; + PyObject *argsbuf[3]; Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 0; PyObject *enabled = Py_True; PyObject *excluding = NULL; + PyObject *eager = NULL; - args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 0, 2, 0, argsbuf); + args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 0, 3, 0, argsbuf); if (!args) { goto exit; } @@ -749,9 +751,15 @@ _imp__set_lazy_imports(PyObject *module, PyObject *const *args, Py_ssize_t nargs if (!noptargs) { goto skip_optional_pos; } - excluding = args[1]; + if (args[1]) { + excluding = args[1]; + if (!--noptargs) { + goto skip_optional_pos; + } + } + eager = args[2]; skip_optional_pos: - return_value = _imp__set_lazy_imports_impl(module, enabled, excluding); + return_value = _imp__set_lazy_imports_impl(module, enabled, excluding, eager); exit: return return_value; @@ -922,4 +930,4 @@ _imp_hydrate_lazy_objects(PyObject *module, PyObject *Py_UNUSED(ignored)) #ifndef _IMP_EXEC_DYNAMIC_METHODDEF #define _IMP_EXEC_DYNAMIC_METHODDEF #endif /* !defined(_IMP_EXEC_DYNAMIC_METHODDEF) */ -/*[clinic end generated code: output=6bee0e7a33fa382e input=a9049054013a1b77]*/ +/*[clinic end generated code: output=cc4f3d25f4d96211 input=a9049054013a1b77]*/ diff --git a/Python/import.c b/Python/import.c index 6c186d5f391..842b221692d 100644 --- 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,7 +4422,7 @@ _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(); @@ -4390,8 +4430,9 @@ PyImport_SetLazyImports(PyObject *enabled, PyObject *excluding) 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,6 +4446,25 @@ 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 { @@ -4412,13 +4472,13 @@ PyImport_SetLazyImports(PyObject *enabled, PyObject *excluding) 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 index 815eb694596..2112f6200d9 100644 --- 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);