Skip to content

Commit cb35c11

Browse files
gh-132775: Add _PyPickle_GetXIData() (gh-133107)
There's some extra complexity due to making sure we we get things right when handling functions and classes defined in the __main__ module. This is also reflected in the tests, including the addition of extra functions in test.support.import_helper.
1 parent 6c522de commit cb35c11

File tree

5 files changed

+1057
-56
lines changed

5 files changed

+1057
-56
lines changed

Include/internal/pycore_crossinterp.h

+7
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,13 @@ PyAPI_FUNC(_PyBytes_data_t *) _PyBytes_GetXIDataWrapped(
171171
xid_newobjfunc,
172172
_PyXIData_t *);
173173

174+
// _PyObject_GetXIData() for pickle
175+
PyAPI_DATA(PyObject *) _PyPickle_LoadFromXIData(_PyXIData_t *);
176+
PyAPI_FUNC(int) _PyPickle_GetXIData(
177+
PyThreadState *,
178+
PyObject *,
179+
_PyXIData_t *);
180+
174181
// _PyObject_GetXIData() for marshal
175182
PyAPI_FUNC(PyObject *) _PyMarshal_ReadObjectFromXIData(_PyXIData_t *);
176183
PyAPI_FUNC(int) _PyMarshal_GetXIData(

Lib/test/support/import_helper.py

+108
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import contextlib
22
import _imp
33
import importlib
4+
import importlib.machinery
45
import importlib.util
56
import os
67
import shutil
@@ -332,3 +333,110 @@ def ensure_lazy_imports(imported_module, modules_to_block):
332333
)
333334
from .script_helper import assert_python_ok
334335
assert_python_ok("-S", "-c", script)
336+
337+
338+
@contextlib.contextmanager
339+
def module_restored(name):
340+
"""A context manager that restores a module to the original state."""
341+
missing = object()
342+
orig = sys.modules.get(name, missing)
343+
if orig is None:
344+
mod = importlib.import_module(name)
345+
else:
346+
mod = type(sys)(name)
347+
mod.__dict__.update(orig.__dict__)
348+
sys.modules[name] = mod
349+
try:
350+
yield mod
351+
finally:
352+
if orig is missing:
353+
sys.modules.pop(name, None)
354+
else:
355+
sys.modules[name] = orig
356+
357+
358+
def create_module(name, loader=None, *, ispkg=False):
359+
"""Return a new, empty module."""
360+
spec = importlib.machinery.ModuleSpec(
361+
name,
362+
loader,
363+
origin='<import_helper>',
364+
is_package=ispkg,
365+
)
366+
return importlib.util.module_from_spec(spec)
367+
368+
369+
def _ensure_module(name, ispkg, addparent, clearnone):
370+
try:
371+
mod = orig = sys.modules[name]
372+
except KeyError:
373+
mod = orig = None
374+
missing = True
375+
else:
376+
missing = False
377+
if mod is not None:
378+
# It was already imported.
379+
return mod, orig, missing
380+
# Otherwise, None means it was explicitly disabled.
381+
382+
assert name != '__main__'
383+
if not missing:
384+
assert orig is None, (name, sys.modules[name])
385+
if not clearnone:
386+
raise ModuleNotFoundError(name)
387+
del sys.modules[name]
388+
# Try normal import, then fall back to adding the module.
389+
try:
390+
mod = importlib.import_module(name)
391+
except ModuleNotFoundError:
392+
if addparent and not clearnone:
393+
addparent = None
394+
mod = _add_module(name, ispkg, addparent)
395+
return mod, orig, missing
396+
397+
398+
def _add_module(spec, ispkg, addparent):
399+
if isinstance(spec, str):
400+
name = spec
401+
mod = create_module(name, ispkg=ispkg)
402+
spec = mod.__spec__
403+
else:
404+
name = spec.name
405+
mod = importlib.util.module_from_spec(spec)
406+
sys.modules[name] = mod
407+
if addparent is not False and spec.parent:
408+
_ensure_module(spec.parent, True, addparent, bool(addparent))
409+
return mod
410+
411+
412+
def add_module(spec, *, parents=True):
413+
"""Return the module after creating it and adding it to sys.modules.
414+
415+
If parents is True then also create any missing parents.
416+
"""
417+
return _add_module(spec, False, parents)
418+
419+
420+
def add_package(spec, *, parents=True):
421+
"""Return the module after creating it and adding it to sys.modules.
422+
423+
If parents is True then also create any missing parents.
424+
"""
425+
return _add_module(spec, True, parents)
426+
427+
428+
def ensure_module_imported(name, *, clearnone=True):
429+
"""Return the corresponding module.
430+
431+
If it was already imported then return that. Otherwise, try
432+
importing it (optionally clear it first if None). If that fails
433+
then create a new empty module.
434+
435+
It can be helpful to combine this with ready_to_import() and/or
436+
isolated_modules().
437+
"""
438+
if sys.modules.get(name) is not None:
439+
mod = sys.modules[name]
440+
else:
441+
mod, _, _ = _force_import(name, False, True, clearnone)
442+
return mod

0 commit comments

Comments
 (0)