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

Lines changed: 7 additions & 0 deletions
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

Lines changed: 108 additions & 0 deletions
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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy