diff --git a/dill/_dill.py b/dill/_dill.py index 668534db..146ab9b6 100644 --- a/dill/_dill.py +++ b/dill/_dill.py @@ -367,7 +367,7 @@ def __init__(self, file, *args, **kwds): self._refonfail = False #settings['dump_module']['refonfail'] if _refonfail is None else _refonfail self._strictio = False #_strictio self._postproc = OrderedDict() - self._file = file # for the logger + self._file_tell = getattr(file, 'tell', None) # for logger and refonfail def dump(self, obj): #NOTE: if settings change, need to update attributes # register if the object is a numpy ufunc @@ -421,32 +421,29 @@ def save(self, obj, save_persistent_id=True, *, name=None): if not self._refonfail: super().save(obj, save_persistent_id) return - if self.framer.current_frame: - # protocol >= 4 - self.framer.commit_frame() - stream = self.framer.current_frame - else: - stream = self._file - position = stream.tell() + # Disable framing (right after the framer.init_framing() call at dump()). + self.framer.current_frame = None + # Store initial state. + position = self._file_tell() memo_size = len(self.memo) try: super().save(obj, save_persistent_id) except UNPICKLEABLE_ERRORS + (AttributeError,) as error_stack: - # AttributeError may happen in save_global() call for child object. + # AttributeError may happen in the save_global() call from a child object. if (type(error_stack) == AttributeError and "no attribute '__name__'" not in error_stack.args[0]): raise - # roll back the stream - stream.seek(position) - stream.truncate() - # roll back memo + # Roll back the stream. + self._file_seek(position) + self._file_truncate() + # Roll back memo. for _ in range(len(self.memo) - memo_size): - self.memo.popitem() # LIFO order is guaranteed for since 3.7 + self.memo.popitem() # LIFO order is guaranteed since 3.7 try: self.save_global(obj, name) except (AttributeError, PicklingError) as error: if getattr(self, '_trace_stack', None) and id(obj) == self._trace_stack[-1]: - # roll back trace state + # Roll back trace state. self._trace_stack.pop() self._size_stack.pop() raise error from error_stack diff --git a/dill/logger.py b/dill/logger.py index 7b6afcdd..385c862d 100644 --- a/dill/logger.py +++ b/dill/logger.py @@ -148,7 +148,7 @@ def trace(self, pickler, msg, *args, obj=None, **kwargs): size = None try: # Streams are not required to be tellable. - size = pickler._file.tell() + size = pickler._file_tell() frame = pickler.framer.current_frame try: size += frame.tell() diff --git a/dill/session.py b/dill/session.py index 9e545f1c..3da31318 100644 --- a/dill/session.py +++ b/dill/session.py @@ -203,10 +203,12 @@ def dump_module( similar but independent from ``dill.settings[`byref`]``, as ``refimported`` refers to virtually all imported objects, while ``byref`` only affects select objects. - refonfail: if `True`, objects that fail to be saved by value will try to - be saved by reference. If it also fails, saving their parent + refonfail: if `True`, objects that fail to pickle by value will try to + be saved by reference. If this also fails, saving their parent objects by reference will be attempted recursively. In the worst - case scenario, the module itself may be saved by reference. + case scenario, the module itself may be saved by reference. Note: + The file-like object must be seekable and truncable with this + option set. **kwds: extra keyword arguments passed to :py:class:`Pickler()`. Raises: @@ -302,9 +304,17 @@ def dump_module( pickler._main = main #FIXME: dill.settings are disabled pickler._byref = False # disable pickling by name reference pickler._recurse = False # disable pickling recursion for globals - pickler._refonfail = refonfail pickler._session = True # is best indicator of when pickling a session pickler._first_pass = True + if refonfail: + pickler._refonfail = True # False by default + pickler._file_seek = getattr(file, 'seek', None) + pickler._file_truncate = getattr(file, 'truncate', None) + if hasattr(file, 'seekable') and not file.seekable(): + pickler._file_seek = None + if pickler._file_seek is None or pickler._file_truncate is None: + raise TypeError("file must have 'tell', 'seek' and 'truncate'" + " attributes if the 'refonfail' option is set.") pickler.dump(main) return