Ticket #15207: richpickle.py

File richpickle.py, 3.0 KB (added by nbruin, 8 years ago)

a possibly light interface

1#by default, python classes are unpickled as an instance with an empty
2#dictionary, which then gets filled in with a __getstate__ call. This
3#construction is guaranteed to work even if there are circular references
4#between objects in the pickle, but when reconstructing circular structures
5#it may happen that the __getstate__ gets executed some time after instantiation
6#and in fact, if the object occurs as a dictionary key, its hash and equality
7#tests may get called prior to __getstate__ being executed.
9#that means that __hash__ etc. may get run with an empty self.__dict__, which
10#probably means it can't execute.
12#the solution is to reinstate the required attributes (but not ones that lead
13#to circularities!) during construction rather than __setstate__. Python
14#offers a protocol for that via __getnewargs__, but implementing that may feel
15#a little heavyweight when all that's necessary is to set some attributes a
16#little earlier. This class tries to offer a little more lightweight version.
17class RichPickle(object):
18    def __reduce__(self):
19        t=type(self)
20        base=RichPickle #we need to find something more appropriate in the mro of t
21        init_attributes = set()
22        for u in t.mro():
23            try:
24                init_attributes.update(u._init_attributes)
25            except AttributeError:
26                pass
27        init_dict=dict()
28        for k in init_attributes:
29            try:
30                init_dict[k]=getattr(self,k)
31            except AttributeError:
32                pass
33        try:
34            rest = self.__getstate__()
35        except AttributeError:
36            try:
37                rest = self.__dict__.copy()
38            except AttributeError:
39                rest = None
40            if rest:
41                for k in init_attributes:
42                    if k in rest:
43                        del rest[k]
44        return (reconstruct_richpickle,(t,base,init_dict),rest)
46#function to help unpickling
47def reconstruct_richpickle(cls, base, init_dict):
48    obj = base.__new__(cls)
49    try:
50        obj.__dict__.update(init_dict)
51    except AttributeError:
52        for key,value in init_dict:
53            setattr(obj,key,value)
54    return obj
56import sage 
57#example class making use of richpickle
58class Node(RichPickle,sage.all.SageObject):
59    def __init__(self,i):
60        self.i=i
61        self.neighbours=set()
63    #setting this attribute indicates that i needs to be set at init
64    #this one line makes circular pickles work again.
65    _init_attributes = ('i',)
67    def __repr__(self):
68        return "Node %s with neighbours %s"%(self.i,sorted([c.i for c in self.neighbours]))
69    def connected_component(self):
70        nodelist=[]
71        todo=[self]
72        visited=set()
73        while todo:
74            n=todo.pop()
75            if n.i not in visited:
76                nodelist.append(n)
77                todo.extend(n.neighbours)
78                visited.add(n.i)
79        nodelist.sort(key=lambda n:n.i)
80        return nodelist
81    def __hash__(self):
82        return self.i