1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
|
# (C) Copyright 2005-2025 Enthought, Inc., Austin, TX
# All rights reserved.
#
# This software is provided without warranty under the terms of the BSD
# license included in LICENSE.txt and may be redistributed only under
# the conditions described in the aforementioned license. The license
# is also available online at http://www.enthought.com/licenses/BSD.txt
#
# Thanks for using Enthought open source!
# Standard library imports
from pickle import _Unpickler as Unpickler
from pickle import UnpicklingError, BUILD
import logging
from types import GeneratorType
# Enthought library imports
from apptools.persistence.updater import __replacement_setstate__
logger = logging.getLogger(__name__)
##############################################################################
# class 'NewUnpickler'
##############################################################################
class NewUnpickler(Unpickler):
"""An unpickler that implements a two-stage pickling process to make it
possible to unpickle complicated Python object hierarchies where the
unserialized state of an object depends on the state of other objects in
the same pickle.
"""
def load(self, max_pass=-1):
"""Read a pickled object representation from the open file.
Return the reconstituted object hierarchy specified in the file.
"""
# List of objects to be unpickled.
self.objects = []
# We overload the load_build method.
dispatch = self.dispatch
dispatch[BUILD[0]] = NewUnpickler.load_build
# call the super class' method.
ret = Unpickler.load(self)
self.initialize(max_pass)
self.objects = []
# Reset the Unpickler's dispatch table.
dispatch[BUILD[0]] = Unpickler.load_build
return ret
def initialize(self, max_pass):
# List of (object, generator) tuples that initialize objects.
generators = []
# Execute object's initialize to setup the generators.
for obj in self.objects:
if hasattr(obj, "__initialize__") and callable(obj.__initialize__):
ret = obj.__initialize__()
if isinstance(ret, GeneratorType):
generators.append((obj, ret))
elif ret is not None:
raise UnpicklingError(
"Unexpected return value from "
"__initialize__. %s returned %s" % (obj, ret)
)
# Ensure a maximum number of passes
if max_pass < 0:
max_pass = len(generators)
# Now run the generators.
count = 0
while len(generators) > 0:
count += 1
if count > max_pass:
not_done = [x[0] for x in generators]
msg = """Reached maximum pass count %s. You may have
a deadlock! The following objects are
uninitialized: %s""" % (
max_pass,
not_done,
)
raise UnpicklingError(msg)
for o, g in generators[:]:
try:
next(g)
except StopIteration:
generators.remove((o, g))
# Make this a class method since dispatch is a class variable.
# Otherwise, supposing the initial VersionedUnpickler.load call (which
# would have overloaded the load_build method) makes a pickle.load call at
# some point, we would have the dispatch still pointing to
# NewPickler.load_build whereas the object being passed in will be an
# Unpickler instance, causing a TypeError.
def load_build(cls, obj):
# Just save the instance in the list of objects.
if isinstance(obj, NewUnpickler):
obj.objects.append(obj.stack[-2])
Unpickler.load_build(obj)
load_build = classmethod(load_build)
class VersionedUnpickler(NewUnpickler):
"""This class reads in a pickled file created at revision version 'n'
and then applies the transforms specified in the updater class to
generate a new set of objects which are at revision version 'n+1'.
I decided to keep the loading of the updater out of this generic class
because we will want updaters to be generated for each plugin's type
of project.
This ensures that the VersionedUnpickler can remain ignorant about the
actual version numbers - all it needs to do is upgrade one release.
"""
def __init__(self, file, updater=None):
Unpickler.__init__(self, file)
self.updater = updater
def find_class(self, module, name):
"""Overridden method from Unpickler.
NB __setstate__ is not called until later.
"""
if self.updater:
# check to see if this class needs to be mapped to a new class
# or module name
original_module, original_name = module, name
module, name = self.updater.get_latest(module, name)
# load the class...
klass = self.import_name(module, name)
# add the updater.... TODO - why the old name?
self.add_updater(original_module, original_name, klass)
else:
# there is no updater so we will be reading in an up to date
# version of the file...
try:
klass = Unpickler.find_class(self, module, name)
except Exception:
logger.error("Looking for [%s] [%s]" % (module, name))
logger.exception(
"Problem using default unpickle functionality"
)
# restore the original __setstate__ if necessary
fn = getattr(klass, "__setstate_original__", False)
if fn:
setattr(klass, "__setstate__", fn)
return klass
def add_updater(self, module, name, klass):
"""If there is an updater defined for this class we will add it to the
class as the __setstate__ method.
"""
fn = self.updater.setstates.get((module, name), False)
if fn:
# move the existing __setstate__ out of the way
self.backup_setstate(module, klass)
# add the updater into the class
setattr(klass, "__updater__", fn)
# hook up our __setstate__ which updates self.__dict__
setattr(klass, "__setstate__", __replacement_setstate__)
else:
pass
def backup_setstate(self, module, klass):
"""If the class has a user defined __setstate__ we back it up."""
if getattr(klass, "__setstate__", False):
if getattr(klass, "__setstate_original__", False):
# don't overwrite the original __setstate__
name = "__setstate__%s" % self.updater.__class__
else:
# backup the original __setstate__ which we will restore
# and run later when we have finished updating the class
name = "__setstate_original__"
method = getattr(klass, "__setstate__")
setattr(klass, name, method)
else:
# the class has no __setstate__ method so do nothing
pass
def import_name(self, module, name):
"""
If the class is needed for the latest version of the application then
it should presumably exist.
If the class no longer exists then we should perhaps return
a proxy of the class.
If the persisted file is at v1 say and the application is at v3 then
objects that are required for v1 and v2 do not have to exist they only
need to be placeholders for the state during an upgrade.
"""
module = __import__(module, globals(), locals(), [name])
return vars(module)[name]
|
