3232from struct import pack , unpack
3333import io
3434import codecs
35+ import weakref
3536import _compat_pickle
3637
3738__all__ = ["PickleError" , "PicklingError" , "UnpicklingError" , "Pickler" ,
@@ -485,6 +486,9 @@ def __init__(self, file, protocol=None, *, fix_imports=True,
485486 self .write = self .framer .write
486487 self ._write_large_bytes = self .framer .write_large_bytes
487488 self .memo = {}
489+ # Ids of objects reached so far only through weak references; any
490+ # left at the end of dump() is an error. See save_weakref.
491+ self ._weak_only = set ()
488492 self .proto = int (protocol )
489493 self .bin = protocol >= 1
490494 self .fast = 0
@@ -499,6 +503,7 @@ def clear_memo(self):
499503 useful when re-using picklers.
500504 """
501505 self .memo .clear ()
506+ self ._weak_only .clear ()
502507
503508 def dump (self , obj ):
504509 """Write a pickled representation of obj to the open file."""
@@ -512,6 +517,14 @@ def dump(self, obj):
512517 if self .proto >= 4 :
513518 self .framer .start_framing ()
514519 self .save (obj )
520+ if self ._weak_only :
521+ # These referents would be collected right after loading. Report
522+ # distinct type names: referents may be numerous and unhashable.
523+ names = {_T (self .memo [i ][1 ])
524+ for i in self ._weak_only if i in self .memo }
525+ raise PicklingError (
526+ "cannot pickle weak reference(s) whose referent(s) have no "
527+ "strong reference in the pickle: " + ", " .join (sorted (names )))
515528 self .write (STOP )
516529 self .framer .end_framing ()
517530
@@ -559,7 +572,7 @@ def get(self, i):
559572
560573 return GET + repr (i ).encode ("ascii" ) + b'\n '
561574
562- def save (self , obj , save_persistent_id = True ):
575+ def save (self , obj , save_persistent_id = True , weak = False ):
563576 self .framer .commit_frame ()
564577
565578 # Check for persistent id (defined by a subclass)
@@ -572,6 +585,9 @@ def save(self, obj, save_persistent_id=True):
572585 # Check the memo
573586 x = self .memo .get (id (obj ))
574587 if x is not None :
588+ # A strong reference anchors the object: no longer weak-only.
589+ if not weak :
590+ self ._weak_only .discard (id (obj ))
575591 self .write (self .get (x [0 ]))
576592 return
577593
@@ -633,6 +649,12 @@ def save(self, obj, save_persistent_id=True):
633649 exc .add_note (f'when serializing { _T (obj )} object' )
634650 raise
635651
652+ # A completed non-weak save anchors obj. Only save_reduce() memoizes
653+ # an object after its arguments, so only here can a weak reference to
654+ # a still-in-progress obj have tagged it.
655+ if not weak :
656+ self ._weak_only .discard (id (obj ))
657+
636658 def persistent_id (self , obj ):
637659 # This exists so a subclass can override it
638660 return None
@@ -1264,6 +1286,40 @@ def save_type(self, obj):
12641286 dispatch [FunctionType ] = save_global
12651287 dispatch [type ] = save_type
12661288
1289+ def save_weakref (self , obj ):
1290+ # Reconstruct as weakref.ref(referent). The referent is saved
1291+ # *weakly* (no anchoring, see save()) and tagged weak-only only after
1292+ # its whole subtree, so a strong cycle inside it cannot anchor it.
1293+ referent = obj ()
1294+ if referent is None :
1295+ # All dead references are interchangeable: pickle by name as the
1296+ # weakref._dead_ref singleton.
1297+ self .write (GLOBAL + b'weakref\n _dead_ref\n ' )
1298+ return
1299+
1300+ save = self .save
1301+ write = self .write
1302+ save (weakref .ref )
1303+ # A referent already in the memo is anchored or already tracked;
1304+ # otherwise it is tagged below, until a strong save of it.
1305+ anchored = id (referent ) in self .memo
1306+ # Build the 1-tuple (referent,); TUPLE1 needs protocol 2.
1307+ if self .proto < 2 :
1308+ write (MARK )
1309+ save (referent , weak = True )
1310+ if not anchored and id (referent ) in self .memo :
1311+ self ._weak_only .add (id (referent ))
1312+ write (TUPLE1 if self .proto >= 2 else TUPLE )
1313+ write (REDUCE )
1314+ # The referent's subtree may have pickled this very (cached) weakref
1315+ # already; reuse the memoized one.
1316+ if id (obj ) in self .memo :
1317+ write (POP + self .get (self .memo [id (obj )][0 ]))
1318+ else :
1319+ self .memoize (obj )
1320+
1321+ dispatch [weakref .ReferenceType ] = save_weakref
1322+
12671323
12681324# Unpickling machinery
12691325
0 commit comments