1"""
2Implementation of NSCoding for OC_PythonObject and friends
3
4NOTE: this only works with a keyed archiver, not with a plain archiver. It
5should be easy enough to change this later on if needed.
6
7A minor problem with NSCoding support is that NSCoding restores
8graphs recusively while Pickle does so depth-first (more of less).
9This can cause problems when the object state contains the
10object itself, which is why we need a 'setValue' callback for the
11load_* functions below.
12"""
13import objc
14from types import *
15import copy_reg
16import copy
17import sys
18
19from pickle import PicklingError, UnpicklingError, whichmodule
20
21
22
23def setupPythonObject():
24    OC_PythonObject = objc.lookUpClass("OC_PythonObject")
25
26    NSArray = objc.lookUpClass("NSArray")
27    NSDictionary = objc.lookUpClass("NSDictionary")
28    NSString = objc.lookUpClass("NSString")
29
30    kOP_REDUCE=0
31    kOP_INST=1
32    kOP_GLOBAL=2
33    kOP_NONE=3
34    kOP_BOOL=4
35    kOP_INT=5
36    kOP_LONG=6
37    kOP_FLOAT=7
38    kOP_UNICODE=8
39    kOP_STRING=9
40    kOP_TUPLE=10
41    kOP_LIST=11
42    kOP_DICT=12
43    kOP_GLOBAL_EXT=13
44    kOP_FLOAT_STR=14
45
46    kKIND = NSString.stringWithString_(u"kind")
47    kFUNC = NSString.stringWithString_(u"func")
48    kARGS = NSString.stringWithString_(u"args")
49    kLIST = NSString.stringWithString_(u"list")
50    kDICT = NSString.stringWithString_(u"dict")
51    kSTATE = NSString.stringWithString_(u"state")
52    kCLASS = NSString.stringWithString_(u"class")
53    kVALUE = NSString.stringWithString_(u"value")
54    kNAME = NSString.stringWithString_(u"name")
55    kMODULE = NSString.stringWithString_(u"module")
56    kCODE = NSString.stringWithString_(u"code")
57
58    class _EmptyClass:
59        pass
60
61    encode_dispatch = {}
62
63    # Code below tries to mirror the implementation in pickle.py, with
64    # adaptations because we're not saving to a byte stream but to another
65    # serializer.
66
67    def save_reduce(coder, func, args,
68            state=None, listitems=None, dictitems=None, obj=None):
69
70        if not isinstance(args, TupleType):
71            raise PicklingError("args from reduce() should be a tuple")
72
73        if not callable(func):
74            raise PicklingError("func from reduce should be callable")
75
76
77        coder.encodeInt_forKey_(kOP_REDUCE, kKIND)
78        coder.encodeObject_forKey_(func, kFUNC)
79        coder.encodeObject_forKey_(args, kARGS)
80        if listitems is None:
81            coder.encodeObject_forKey_(None, kLIST)
82        else:
83            coder.encodeObject_forKey_(list(listitems), kLIST)
84
85        if dictitems is None:
86            coder.encodeObject_forKey_(None, kDICT)
87        else:
88            coder.encodeObject_forKey_(dict(dictitems), kDICT)
89        coder.encodeObject_forKey_(state, kSTATE)
90
91    def save_inst(coder, obj):
92        if hasattr(obj, '__getinitargs__'):
93            args = obj.__getinitargs__()
94            len(args) # Assert it's a sequence
95        else:
96            args = ()
97
98        cls = obj.__class__
99
100        coder.encodeInt32_forKey_(kOP_INST, kKIND)
101        coder.encodeObject_forKey_(cls, kCLASS)
102        coder.encodeObject_forKey_(args, kARGS)
103
104        try:
105            getstate = obj.__getstate__
106        except AttributeError:
107            state = obj.__dict__
108
109        else:
110            state = getstate()
111
112        coder.encodeObject_forKey_(state, kSTATE)
113
114    encode_dispatch[InstanceType] = save_inst
115
116
117    def save_none(coder, obj):
118        coder.encodeInt_forKey_(kOP_NONE, kKIND)
119    encode_dispatch[NoneType] = save_none
120
121    def save_bool(coder, obj):
122        coder.encodeInt_forKey_(kOP_BOOL, kKIND)
123        coder.encodeBool_forKey_(bool(obj), kVALUE)
124    encode_dispatch[bool] = save_bool
125
126    def save_int(coder, obj):
127        coder.encodeInt_forKey_(kOP_INT, kKIND)
128        coder.encodeInt64_forKey_(obj, kVALUE)
129    encode_dispatch[int] = save_int
130
131    def save_long(coder, obj):
132        coder.encodeInt_forKey_(kOP_LONG, kKIND)
133        coder.encodeObject_forKey_(unicode(repr(obj)), kVALUE)
134    encode_dispatch[long] = save_long
135
136    def save_float(coder, obj):
137        # Encode floats as strings, this seems to be needed to get
138        # 100% reliable round-trips.
139        coder.encodeInt_forKey_(kOP_FLOAT_STR, kKIND)
140        coder.encodeObject_forKey_(unicode(repr(obj)), kVALUE)
141        #coder.encodeDouble_forKey_(obj, kVALUE)
142    encode_dispatch[float] = save_float
143
144    def save_string(coder, obj):
145        coder.encodeInt_forKey_(kOP_STRING, kKIND)
146        coder.encodeBytes_length_forKey_(obj, len(obj), kVALUE)
147    encode_dispatch[str] = save_string
148
149
150    def save_tuple(coder, obj):
151        coder.encodeInt_forKey_(kOP_TUPLE, kKIND)
152        coder.encodeObject_forKey_(NSArray.arrayWithArray_(obj), kVALUE)
153    encode_dispatch[tuple] = save_tuple
154
155    def save_list(coder, obj):
156        coder.encodeInt_forKey_(kOP_LIST, kKIND)
157        coder.encodeObject_forKey_(NSArray.arrayWithArray_(obj), kVALUE)
158    encode_dispatch[list] = save_list
159
160    def save_dict(coder, obj):
161        coder.encodeInt_forKey_(kOP_DICT, kKIND)
162        v = NSDictionary.dictionaryWithDictionary_(obj)
163        coder.encodeObject_forKey_(v, kVALUE)
164    encode_dispatch[dict] = save_dict
165
166    def save_global(coder, obj, name=None):
167
168        if name is None:
169            name = obj.__name__
170
171        module = getattr(obj, "__module__", None)
172        if module is None:
173            module = whichmodule(obj, name)
174
175        try:
176            __import__ (module)
177            mod = sys.modules[module]
178            klass= getattr(mod, name)
179
180        except (ImportError, KeyError, AttributeError):
181            raise PicklingError(
182                      "Can't pickle %r: it's not found as %s.%s" %
183                      (obj, module, name))
184        else:
185            if klass is not obj:
186                raise PicklingError(
187                    "Can't pickle %r: it's not the same object as %s.%s" %
188                    (obj, module, name))
189
190        code = copy_reg._extension_registry.get((module, name))
191        if code:
192            coder.encodeInt_forKey_(kOP_GLOBAL_EXT, kKIND)
193            coder.encodeInt_forKey_(code, kCODE)
194
195        else:
196            coder.encodeInt_forKey_(kOP_GLOBAL, kKIND)
197            coder.encodeObject_forKey_(unicode(module), kMODULE)
198            coder.encodeObject_forKey_(unicode(name), kNAME)
199
200    encode_dispatch[ClassType] = save_global
201    encode_dispatch[FunctionType] = save_global
202    encode_dispatch[BuiltinFunctionType] = save_global
203    encode_dispatch[TypeType] = save_global
204
205
206    decode_dispatch = {}
207
208    def load_none(coder, setValue):
209        return None
210    decode_dispatch[kOP_NONE] = load_none
211
212    def load_bool(coder, setValue):
213        return coder.decodeBoolForKey_(kVALUE)
214    decode_dispatch[kOP_BOOL] = load_bool
215
216    def load_int(coder, setValue):
217        return int(coder.decodeInt64ForKey_(kVALUE))
218    decode_dispatch[kOP_INT] = load_int
219
220    def load_long(coder, setValue):
221        return long(coder.decodeObjectForKey_(kVALUE))
222    decode_dispatch[kOP_LONG] = load_long
223
224    def load_float(coder, setValue):
225        return coder.decodeFloatForKey_(kVALUE)
226    decode_dispatch[kOP_FLOAT] = load_float
227
228    def load_float_str(coder, setValue):
229        return float(coder.decodeObjectForKey_(kVALUE))
230    decode_dispatch[kOP_FLOAT_STR] = load_float_str
231
232    def load_tuple(coder, setValue):
233        return tuple(coder.decodeObjectForKey_(kVALUE))
234    decode_dispatch[kOP_TUPLE] = load_tuple
235
236    def load_list(coder, setValue):
237        return list(coder.decodeObjectForKey_(kVALUE))
238    decode_dispatch[kOP_LIST] = load_list
239
240    def load_dict(coder, setValue):
241        return dict(coder.decodeObjectForKey_(kVALUE))
242    decode_dispatch[kOP_DICT] = load_dict
243
244    def load_global_ext(coder, setValue):
245        code = coder.intForKey_(kCODE)
246        nil = []
247        obj = copy_reg._extension_cache.get(code, nil)
248        if obj is not nil:
249            return obj
250        key = copy_reg._inverted_registry.get(code)
251        if not key:
252            raise ValueError("unregistered extension code %d" % code)
253
254        module, name = key
255        __import__(module)
256        mod = sys.modules[module]
257        klass = getattr(mod, name)
258        return klass
259    decode_dispatch[kOP_GLOBAL_EXT] = load_global_ext
260
261    def load_global(coder, setValue):
262        module = coder.decodeObjectForKey_(kMODULE)
263        name = coder.decodeObjectForKey_(kNAME)
264        __import__(module)
265        mod = sys.modules[module]
266        klass = getattr(mod, name)
267        return klass
268
269    decode_dispatch[kOP_GLOBAL] = load_global
270
271
272    def load_inst(coder, setValue):
273        cls = coder.decodeObjectForKey_(kCLASS)
274        initargs = coder.decodeObjectForKey_(kARGS)
275
276        instantiated = 0
277        if (not initargs and
278                type(cls) is ClassType and
279                not hasattr(cls, "__getinitargs__")):
280            try:
281                value = _EmptyClass()
282                value.__class__ = cls
283                instantiated = 1
284
285            except RuntimeError:
286                pass
287
288        if not instantiated:
289            try:
290                value = cls(*initargs)
291            except TypeError, err:
292                raise TypeError, "in constructor for %s: %s" % (
293                    cls.__name__, str(err)), sys.exc_info()[2]
294
295
296        # We now have the object, but haven't set the correct
297        # state yet.  Tell the bridge about this value right
298        # away, that's needed because `value` might be part
299        # of the object state which we'll retrieve next.
300        setValue(value)
301
302        state = coder.decodeObjectForKey_(kSTATE)
303        setstate = getattr(value, "__setstate__", None)
304        if setstate is not None:
305            setstate(state)
306            return value
307
308        slotstate = None
309        if isinstance(state, tuple) and len(state) == 2:
310            state, slotstate = state
311
312        if state:
313            try:
314                value.__dict__.update(state)
315            except RuntimeError:
316                for k, v in state.items():
317                    setattr(value, k, v)
318
319        if slotstate:
320            for k, v in slotstate.items():
321                setattr(value, k, v)
322
323        return value
324    decode_dispatch[kOP_INST] = load_inst
325
326
327    def load_reduce(coder, setValue):
328        func = coder.decodeObjectForKey_(kFUNC)
329        args = coder.decodeObjectForKey_(kARGS)
330
331        value = func(*args)
332
333        # We now have the object, but haven't set the correct
334        # state yet.  Tell the bridge about this value right
335        # away, that's needed because `value` might be part
336        # of the object state which we'll retrieve next.
337        setValue(value)
338
339        listitems = coder.decodeObjectForKey_(kLIST)
340        dictitems = coder.decodeObjectForKey_(kDICT)
341        state = coder.decodeObjectForKey_(kSTATE)
342        setstate = getattr(value, "__setstate__", None)
343        if setstate:
344            setstate(state)
345            return
346
347        slotstate = None
348        if isinstance(state, tuple) and len(state) == 2:
349            state, slotstate = state
350
351        if state:
352            try:
353                value.__dict__.update(state)
354
355            except RuntimeError:
356                for k, v in state.items():
357                    setattr(value, k, v)
358
359        if slotstate:
360            for k, v in slotstate.items():
361                setattr(value, k, v)
362
363        if listitems:
364            for a in listitems:
365                value.append(a)
366
367        if dictitems:
368            for k, v in dictitems.items():
369                value[k] = v
370
371        return value
372    decode_dispatch[kOP_REDUCE] = load_reduce
373
374
375    def pyobjectEncode(self, coder):
376        t = type(self)
377
378        # Find builtin support
379        f = encode_dispatch.get(t)
380        if f is not None:
381            f(coder, self)
382            return
383
384        # Check for a class with a custom metaclass
385        try:
386            issc = issubclass(t, TypeType)
387        except TypeError:
388            issc = 0
389
390        if issc:
391            save_global(coder, self)
392            return
393
394        # Check copy_reg.dispatch_table
395        reduce = copy_reg.dispatch_table.get(t)
396        if reduce is not None:
397            rv = reduce(self)
398
399        else:
400            reduce = getattr(self, "__reduce_ex__", None)
401            if reduce is not None:
402                rv = reduce(2)
403
404            else:
405                rv = getattr(self, "__reduce__", None)
406                if reduce is not None:
407                    rv = reduce()
408
409                else:
410                    raise PicklingError("Can't pickle %r object: %r" %
411                            (t.__name__, self))
412
413        if type(rv) is StringType:
414            save_global(coder, rv)
415            return
416
417        if type(rv) is not TupleType:
418            raise PicklingError("%s must return string or tuple" % reduce)
419
420        l = len(rv)
421        if not (2 <= l <= 5):
422            raise PicklingError("Tuple returned by %s must have two to "
423                    "five elements" % reduce)
424
425        save_reduce(coder, *rv)
426
427    def pyobjectDecode(coder, setValue):
428        tp = coder.decodeIntForKey_(kKIND)
429        f = decode_dispatch.get(tp)
430        if f is None:
431            raise UnpicklingError("Unknown object kind: %s"%(tp,))
432
433        return f(coder, setValue)
434
435    # An finally register the coder/decoder
436    OC_PythonObject.setVersion_coder_decoder_copier_(
437            1, pyobjectEncode, pyobjectDecode, copy.copy)
438
439setupPythonObject()
440