diff --git a/dill/_dill.py b/dill/_dill.py index 0130e709..a9662522 100644 --- a/dill/_dill.py +++ b/dill/_dill.py @@ -229,10 +229,9 @@ def dump(obj, file, protocol=None, byref=None, fmode=None, recurse=None, **kwds) See :func:`dumps` for keyword arguments. """ from .settings import settings - protocol = settings['protocol'] if protocol is None else int(protocol) - _kwds = kwds.copy() - _kwds.update(dict(byref=byref, fmode=fmode, recurse=recurse)) - Pickler(file, protocol, **_kwds).dump(obj) + protocol = int(_getopt(settings, 'protocol', protocol)) + kwds.update(byref=byref, fmode=fmode, recurse=recurse) + Pickler(file, protocol, **kwds).dump(obj) return def dumps(obj, protocol=None, byref=None, fmode=None, recurse=None, **kwds):#, strictio=None): @@ -317,6 +316,26 @@ class PicklingWarning(PickleWarning, PicklingError): class UnpicklingWarning(PickleWarning, UnpicklingError): pass +def _getopt(settings, key, arg=None, *, kwds=None): + """Get option from named argument 'arg' or 'kwds', falling back to settings. + + Examples: + + # With an explicitly named argument: + protocol = int(_getopt(settings, 'protocol', protocol)) + + # With a named argument in **kwds: + self._byref = _getopt(settings, 'byref', kwds=kwds) + """ + # Sanity check, it's a bug in calling code if False. + assert kwds is None or arg is None + if kwds is not None: + arg = kwds.pop(key, None) + if arg is not None: + return arg + else: + return settings[key] + ### Extend the Picklers class Pickler(StockPickler): """python's Pickler extended to interpreter sessions""" @@ -326,19 +345,15 @@ class Pickler(StockPickler): def __init__(self, file, *args, **kwds): settings = Pickler.settings - _byref = kwds.pop('byref', None) - #_strictio = kwds.pop('strictio', None) - _fmode = kwds.pop('fmode', None) - _recurse = kwds.pop('recurse', None) - StockPickler.__init__(self, file, *args, **kwds) self._main = _main_module self._diff_cache = {} - self._byref = settings['byref'] if _byref is None else _byref - self._strictio = False #_strictio - self._fmode = settings['fmode'] if _fmode is None else _fmode - self._recurse = settings['recurse'] if _recurse is None else _recurse + self._byref = _getopt(settings, 'byref', kwds=kwds) + self._fmode = _getopt(settings, 'fmode', kwds=kwds) + self._recurse = _getopt(settings, 'recurse', kwds=kwds) + self._strictio = False #_getopt(settings, 'strictio', kwds=kwds) self._postproc = OrderedDict() self._file = file + StockPickler.__init__(self, file, *args, **kwds) def save(self, obj, save_persistent_id=True): # register if the object is a numpy ufunc @@ -410,10 +425,9 @@ def find_class(self, module, name): def __init__(self, *args, **kwds): settings = Pickler.settings - _ignore = kwds.pop('ignore', None) - StockUnpickler.__init__(self, *args, **kwds) self._main = _main_module - self._ignore = settings['ignore'] if _ignore is None else _ignore + self._ignore = _getopt(settings, 'ignore', kwds=kwds) + StockUnpickler.__init__(self, *args, **kwds) def load(self): #NOTE: if settings change, need to update attributes obj = StockUnpickler.load(self) diff --git a/dill/session.py b/dill/session.py index 6acdd432..031fe428 100644 --- a/dill/session.py +++ b/dill/session.py @@ -22,8 +22,8 @@ from dill import _dill, Pickler, Unpickler from ._dill import ( BuiltinMethodType, FunctionType, MethodType, ModuleType, TypeType, - _import_module, _is_builtin_module, _is_imported_module, _main_module, - _reverse_typemap, __builtin__, + _getopt, _import_module, _is_builtin_module, _is_imported_module, + _main_module, _reverse_typemap, __builtin__, ) # Type hints. @@ -211,8 +211,10 @@ def dump_module( refimported = kwds.pop('byref', refimported) module = kwds.pop('main', module) - from .settings import settings - protocol = settings['protocol'] + from .settings import settings as dill_settings + protocol = dill_settings['protocol'] + refimported = _getopt(settings, 'refimported', refimported) + main = module if main is None: main = _main_module @@ -576,6 +578,12 @@ def load_module_asdict( main.__session__ = str(filename) return main.__dict__ +## Session settings ## + +settings = { + 'refimported': False, +} + # Internal exports for backward compatibility with dill v0.3.5.1 # Can't be placed in dill._dill because of circular import problems.