|
5 | 5 |
|
6 | 6 |
|
7 | 7 | import zlib |
8 | | -import re |
9 | 8 | import cPickle as pickle_module |
10 | | -import pickle # Importing just to get dispatch table, not pickling with it. |
11 | | -import copy_reg |
12 | | -import types |
13 | | - |
14 | | -from python_toolbox import address_tools |
15 | | -from python_toolbox import misc_tools |
16 | | - |
17 | | - |
18 | | -def is_atomically_pickleable(thing): |
19 | | - ''' |
20 | | - Return whether `thing` is an atomically pickleable object. |
21 | | - |
22 | | - "Atomically-pickleable" means that it's pickleable without considering any |
23 | | - other object that it contains or refers to. For example, a `list` is |
24 | | - atomically pickleable, even if it contains an unpickleable object, like a |
25 | | - `threading.Lock()`. |
26 | | - |
27 | | - However, the `threading.Lock()` itself is not atomically pickleable. |
28 | | - ''' |
29 | | - my_type = misc_tools.get_actual_type(thing) |
30 | | - return _is_type_atomically_pickleable(my_type, thing) |
31 | | - |
32 | | - |
33 | | -def _is_type_atomically_pickleable(type_, thing=None): |
34 | | - '''Return whether `type_` is an atomically pickleable type.''' |
35 | | - try: |
36 | | - return _is_type_atomically_pickleable.cache[type_] |
37 | | - except KeyError: |
38 | | - pass |
39 | | - |
40 | | - if thing is not None: |
41 | | - assert isinstance(thing, type_) |
42 | | - |
43 | | - # Sub-function in order to do caching without crowding the main algorithm: |
44 | | - def get_result(): |
45 | | - |
46 | | - # We allow a flag for types to painlessly declare whether they're |
47 | | - # atomically pickleable: |
48 | | - if hasattr(type_, '_is_atomically_pickleable'): |
49 | | - return type_._is_atomically_pickleable |
50 | | - |
51 | | - # Weird special case: `threading.Lock` objects don't have `__class__`. |
52 | | - # We assume that objects that don't have `__class__` can't be pickled. |
53 | | - # (With the exception of old-style classes themselves.) |
54 | | - if not hasattr(thing, '__class__') and \ |
55 | | - (not isinstance(thing, types.ClassType)): |
56 | | - return False |
57 | | - |
58 | | - if not issubclass(type_, object): |
59 | | - return True |
60 | | - |
61 | | - def assert_legit_pickling_exception(exception): |
62 | | - '''Assert that `exception` reports a problem in pickling.''' |
63 | | - message = exception.args[0] |
64 | | - segments = [ |
65 | | - "can't pickle", |
66 | | - 'should only be shared between processes through inheritance', |
67 | | - 'cannot be passed between processes or pickled' |
68 | | - ] |
69 | | - assert any((segment in message) for segment in segments) |
70 | | - # todo: turn to warning |
71 | | - |
72 | | - if type_ in pickle.Pickler.dispatch: |
73 | | - return True |
74 | | - |
75 | | - reduce_function = copy_reg.dispatch_table.get(type_) |
76 | | - if reduce_function: |
77 | | - try: |
78 | | - reduce_result = reduce_function(thing) |
79 | | - except Exception, exception: |
80 | | - assert_legit_pickling_exception(exception) |
81 | | - return False |
82 | | - else: |
83 | | - return True |
84 | | - |
85 | | - reduce_function = getattr(type_, '__reduce_ex__', None) |
86 | | - if reduce_function: |
87 | | - try: |
88 | | - reduce_result = reduce_function(thing, 0) |
89 | | - # (The `0` is the protocol argument.) |
90 | | - except Exception, exception: |
91 | | - assert_legit_pickling_exception(exception) |
92 | | - return False |
93 | | - else: |
94 | | - return True |
95 | | - |
96 | | - reduce_function = getattr(type_, '__reduce__', None) |
97 | | - if reduce_function: |
98 | | - try: |
99 | | - reduce_result = reduce_function(thing) |
100 | | - except Exception, exception: |
101 | | - assert_legit_pickling_exception(exception) |
102 | | - return False |
103 | | - else: |
104 | | - return True |
105 | | - |
106 | | - return False |
107 | | - |
108 | | - result = get_result() |
109 | | - _is_type_atomically_pickleable.cache[type_] = result |
110 | | - return result |
111 | | - |
112 | | -_is_type_atomically_pickleable.cache = {} |
113 | | - |
114 | | - |
115 | | -class FilteredObject(object): |
116 | | - '''Placeholder for an object that was filtered out when pickling.''' |
117 | | - def __init__(self, about): |
118 | | - self.about = about |
119 | | - def __repr__(self): |
120 | | - return 'FilteredObject(%s)' % repr(self.about) |
121 | | - def __getattr__(self, key): |
122 | | - return FilteredObject('%s.%s' % (self.about, key)) |
123 | | - |
124 | | - |
125 | | -_filtered_string_pattern = re.compile( |
126 | | - r'^Filtered by pickle_tools \((?P<description>.*?)\)$' |
127 | | -) |
128 | | - |
129 | | -class CutePickler(object): |
130 | | - ''' |
131 | | - Pickler which filters out non-pickleable objects. |
132 | | - |
133 | | - When the pickler comes upon a non-pickleable object it replaces it with a |
134 | | - marker which will cause it to become a `FilteredObject` upon unpickling. |
135 | | - |
136 | | - (Not subclassing `cPickle.Pickler` because it doesn't support subclassing.) |
137 | | - ''' |
138 | | - def __init__(self, file_, protocol=0): |
139 | | - pickler = self.pickler = pickle_module.Pickler(file_, protocol) |
140 | | - pickler.persistent_id = self.persistent_id |
141 | | - self.dump, self.clear_memo = \ |
142 | | - pickler.dump, pickler.clear_memo |
143 | | - |
144 | | - |
145 | | - def persistent_id(self, obj): |
146 | | - if self.pre_filter: |
147 | | - passed_pre_filter = self.pre_filter(obj) |
148 | | - else: |
149 | | - passed_pre_filter = True |
150 | | - |
151 | | - if passed_pre_filter and is_atomically_pickleable(obj): |
152 | | - return None |
153 | | - else: |
154 | | - return 'Filtered by pickle_tools (%s)' % \ |
155 | | - address_tools.describe(obj) |
156 | | - |
157 | | - def pre_filter(self, thing): |
158 | | - '''Pre-filter `thing`, returning `False` if it shouldn't be pickled.''' |
159 | | - return True |
160 | | - |
161 | | - |
162 | | -class CuteUnpickler(object): |
163 | | - ''' |
164 | | - Unpickler which replaces non-pickleable objects with `FilteredObject`s. |
165 | | - |
166 | | - When the corresponding `CutePickler` came upon non-pickleable objects it |
167 | | - replacef them with a markers which will cause `CuteUnpickler` to replace |
168 | | - them with `FilteredObject` instances upon unpickling. |
169 | | - |
170 | | - (Not subclassing `cPickle.Unpickler` because it doesn't support |
171 | | - subclassing.) |
172 | | - ''' |
173 | | - def __init__(self, file_): |
174 | | - unpickler = self.unpickler = pickle_module.Unpickler(file_) |
175 | | - unpickler.persistent_load = self.persistent_load |
176 | | - self.load = unpickler.load |
177 | | - self.noload = getattr(unpickler, 'noload', None) |
178 | | - # (Defaulting to `None` because `pickle.Unpickler` doesn't have |
179 | | - # `noload`.) |
180 | | - |
181 | | - def persistent_load(self, id_string): |
182 | | - match = _filtered_string_pattern.match(id_string) |
183 | | - if match: |
184 | | - description = match.groupdict()['description'] |
185 | | - return FilteredObject(description) |
186 | | - else: |
187 | | - raise pickle_module.UnpicklingError('Invalid persistent id') |
188 | 9 |
|
189 | 10 |
|
190 | 11 | def compickle(thing): |
|
0 commit comments