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 sys
14import objc
15from types import *
16import copy_reg
17import copy
18
19from pickle import PicklingError, UnpicklingError, whichmodule
20
21if hasattr(sys, 'intern'):
22    intern = sys.intern
23
24
25# FIXME: This requires manual wrappers from the Foundation bindings
26def setupPythonObject():
27    OC_PythonObject = objc.lookUpClass("OC_PythonObject")
28
29    NSArray = objc.lookUpClass("NSArray")
30    NSDictionary = objc.lookUpClass("NSDictionary")
31    NSString = objc.lookUpClass("NSString")
32
33    kOP_REDUCE=0
34    kOP_INST=1
35    kOP_GLOBAL=2
36    kOP_NONE=3
37    kOP_BOOL=4
38    kOP_INT=5
39    kOP_LONG=6
40    kOP_FLOAT=7
41    kOP_UNICODE=8
42    kOP_STRING=9
43    kOP_TUPLE=10
44    kOP_LIST=11
45    kOP_DICT=12
46    kOP_GLOBAL_EXT=13
47    kOP_FLOAT_STR=14
48
49    kKIND = NSString.stringWithString_(u"kind")
50    kFUNC = NSString.stringWithString_(u"func")
51    kARGS = NSString.stringWithString_(u"args")
52    kLIST = NSString.stringWithString_(u"list")
53    kDICT = NSString.stringWithString_(u"dict")
54    kSTATE = NSString.stringWithString_(u"state")
55    kCLASS = NSString.stringWithString_(u"class")
56    kVALUE = NSString.stringWithString_(u"value")
57    kNAME = NSString.stringWithString_(u"name")
58    kMODULE = NSString.stringWithString_(u"module")
59    kCODE = NSString.stringWithString_(u"code")
60
61    class _EmptyClass:
62        pass
63
64    encode_dispatch = {}
65
66    # Code below tries to mirror the implementation in pickle.py, with
67    # adaptations because we're not saving to a byte stream but to another
68    # serializer.
69
70    def save_reduce(coder, func, args,
71            state=None, listitems=None, dictitems=None, obj=None):
72
73        if not isinstance(args, tuple):
74            raise PicklingError("args from reduce() should be a tuple")
75
76        if not callable(func):
77            raise PicklingError("func from reduce should be callable")
78
79        if coder.allowsKeyedCoding():
80            coder.encodeInt_forKey_(kOP_REDUCE, kKIND)
81            coder.encodeObject_forKey_(func, kFUNC)
82            coder.encodeObject_forKey_(args, kARGS)
83            if listitems is None:
84                coder.encodeObject_forKey_(None, kLIST)
85            else:
86                coder.encodeObject_forKey_(list(listitems), kLIST)
87
88            if dictitems is None:
89                coder.encodeObject_forKey_(None, kDICT)
90            else:
91                coder.encodeObject_forKey_(dict(dictitems), kDICT)
92            coder.encodeObject_forKey_(state, kSTATE)
93
94        else:
95            coder.__pyobjc__encodeInt_(kOP_REDUCE)
96            coder.encodeObject_(func)
97            coder.encodeObject_(args)
98            if listitems is None:
99                coder.encodeObject_(None)
100            else:
101                coder.encodeObject_(list(listitems))
102
103            if dictitems is None:
104                coder.encodeObject_(None)
105            else:
106                coder.encodeObject_(dict(dictitems))
107            coder.encodeObject_(state)
108
109    if sys.version_info[0] == 2:
110        def save_inst(coder, obj):
111            if hasattr(obj, '__getinitargs__'):
112                args = obj.__getinitargs__()
113                len(args) # Assert it's a sequence
114            else:
115                args = ()
116
117            cls = obj.__class__
118
119            if coder.allowsKeyedCoding():
120                coder.encodeInt32_forKey_(kOP_INST, kKIND)
121                coder.encodeObject_forKey_(cls, kCLASS)
122                coder.encodeObject_forKey_(args, kARGS)
123
124            else:
125                coder.__pyobjc__encodeInt32_(kOP_INST)
126                coder.encodeObject_(cls)
127                coder.encodeObject_(args)
128
129            try:
130                getstate = obj.__getstate__
131            except AttributeError:
132                state = obj.__dict__
133
134            else:
135                state = getstate()
136
137            if coder.allowsKeyedCoding():
138                coder.encodeObject_forKey_(state, kSTATE)
139
140            else:
141                coder.encodeObject_(state)
142
143        encode_dispatch[InstanceType] = save_inst
144
145
146    def save_none(coder, obj):
147        if coder.allowsKeyedCoding():
148            coder.encodeInt_forKey_(kOP_NONE, kKIND)
149        else:
150            coder.__pyobjc__encodeInt_(kOP_NONE)
151    encode_dispatch[type(None)] = save_none
152
153    def save_bool(coder, obj):
154        if coder.allowsKeyedCoding():
155            coder.encodeInt_forKey_(kOP_BOOL, kKIND)
156            coder.encodeBool_forKey_(bool(obj), kVALUE)
157        else:
158            coder.__pyobjc__encodeInt_(kOP_BOOL)
159            coder.__pyobjc__encodeBool_(bool(obj))
160    encode_dispatch[bool] = save_bool
161
162    if sys.version_info[0] == 2:
163        def save_int(coder, obj):
164            if coder.allowsKeyedCoding():
165                coder.encodeInt_forKey_(kOP_INT, kKIND)
166                coder.encodeInt64_forKey_(obj, kVALUE)
167            else:
168                coder.__pyobjc__encodeInt_(kOP_INT)
169                coder.__pyobjc__encodeInt64_(obj)
170        encode_dispatch[int] = save_int
171
172        def save_long(coder, obj):
173            if coder.allowsKeyedCoding():
174                coder.encodeInt_forKey_(kOP_LONG, kKIND)
175                coder.encodeObject_forKey_(unicode(repr(obj)), kVALUE)
176            else:
177                coder.__pyobjc__encodeInt_(kOP_LONG)
178                coder.encodeObject_(unicode(repr(obj)))
179
180        encode_dispatch[long] = save_long
181
182    else:
183        def save_int(coder, obj):
184            if coder.allowsKeyedCoding():
185                coder.encodeInt_forKey_(kOP_LONG, kKIND)
186                coder.encodeObject_forKey_(unicode(repr(obj)), kVALUE)
187            else:
188                coder.__pyobjc__encodeInt_(kOP_LONG)
189                coder.encodeObject_(unicode(repr(obj)))
190        encode_dispatch[int] = save_int
191
192    def save_float(coder, obj):
193        # Encode floats as strings, this seems to be needed to get
194        # 100% reliable round-trips.
195        if coder.allowsKeyedCoding():
196            coder.encodeInt_forKey_(kOP_FLOAT_STR, kKIND)
197            coder.encodeObject_forKey_(unicode(repr(obj)), kVALUE)
198        else:
199            coder.__pyobjc__encodeInt_(kOP_FLOAT_STR)
200            coder.encodeObject_(unicode(repr(obj)))
201        #coder.encodeDouble_forKey_(obj, kVALUE)
202    encode_dispatch[float] = save_float
203
204    def save_string(coder, obj):
205        if coder.allowsKeyedCoding():
206            coder.encodeInt_forKey_(kOP_STRING, kKIND)
207            coder.encodeBytes_length_forKey_(obj, len(obj), kVALUE)
208        else:
209            encodeInt_(kOP_STRING)
210            coder.encodeBytes_length_(obj, len(obj))
211
212    encode_dispatch[str] = save_string
213
214
215    def save_tuple(coder, obj):
216        if coder.allowsKeyedCoding():
217            coder.encodeInt_forKey_(kOP_TUPLE, kKIND)
218            coder.encodeObject_forKey_(NSArray.arrayWithArray_(obj), kVALUE)
219
220        else:
221            coder.__pyobjc__encodeInt_(kOP_TUPLE)
222            coder.encodeObject_(NSArray.arrayWithArray_(obj))
223    encode_dispatch[tuple] = save_tuple
224
225    def save_list(coder, obj):
226        if coder.allowsKeyedCoding():
227            coder.encodeInt_forKey_(kOP_LIST, kKIND)
228            coder.encodeObject_forKey_(NSArray.arrayWithArray_(obj), kVALUE)
229
230        else:
231            coder.__pyobjc__encodeInt_(kOP_LIST)
232            coder.encodeObject_(NSArray.arrayWithArray_(obj))
233    encode_dispatch[list] = save_list
234
235    def save_dict(coder, obj):
236        if coder.allowsKeyedCoding():
237            coder.encodeInt_forKey_(kOP_DICT, kKIND)
238            v = NSDictionary.dictionaryWithDictionary_(obj)
239            coder.encodeObject_forKey_(v, kVALUE)
240        else:
241            coder.__pyobjc__encodeInt_(kOP_DICT)
242            v = NSDictionary.dictionaryWithDictionary_(obj)
243            coder.encodeObject_(v)
244
245    encode_dispatch[dict] = save_dict
246
247    def save_global(coder, obj, name=None):
248
249        if name is None:
250            name = obj.__name__
251
252        module = getattr(obj, "__module__", None)
253        if module is None:
254            module = whichmodule(obj, name)
255
256        try:
257            __import__ (module)
258            mod = sys.modules[module]
259            klass= getattr(mod, name)
260
261        except (ImportError, KeyError, AttributeError):
262            raise PicklingError(
263                      "Can't pickle %r: it's not found as %s.%s" %
264                      (obj, module, name))
265        else:
266            if klass is not obj:
267                raise PicklingError(
268                    "Can't pickle %r: it's not the same object as %s.%s" %
269                    (obj, module, name))
270
271        code = copy_reg._extension_registry.get((module, name))
272
273        if coder.allowsKeyedCoding():
274            if code:
275                coder.encodeInt_forKey_(kOP_GLOBAL_EXT, kKIND)
276                coder.encodeInt_forKey_(code, kCODE)
277
278            else:
279                coder.encodeInt_forKey_(kOP_GLOBAL, kKIND)
280                coder.encodeObject_forKey_(unicode(module), kMODULE)
281                coder.encodeObject_forKey_(unicode(name), kNAME)
282
283        else:
284            if code:
285                coder.__pyobjc__encodeInt_(kOP_GLOBAL_EXT)
286                coder.__pyobjc__encodeInt_(code)
287
288            else:
289                coder.__pyobjc__encodeInt_(kOP_GLOBAL)
290                coder.encodeObject_(unicode(module))
291                coder.encodeObject_(unicode(name))
292
293    if sys.version_info[0] == 2:
294        encode_dispatch[ClassType] = save_global
295    encode_dispatch[type(save_global)] = save_global
296    encode_dispatch[type(dir)] = save_global
297    encode_dispatch[type] = save_global
298
299
300    decode_dispatch = {}
301
302    def load_none(coder, setValue):
303        return None
304    decode_dispatch[kOP_NONE] = load_none
305
306    def load_bool(coder, setValue):
307        if coder.allowsKeyedCoding():
308            return coder.decodeBoolForKey_(kVALUE)
309        else:
310            return coder.__pyobjc__decodeBool()
311
312    decode_dispatch[kOP_BOOL] = load_bool
313
314    def load_int(coder, setValue):
315        if coder.allowsKeyedCoding():
316            return int(coder.decodeInt64ForKey_(kVALUE))
317        else:
318            return int(coder.__pyobjc__decodeInt64())
319    decode_dispatch[kOP_INT] = load_int
320
321    def load_long(coder, setValue):
322        if coder.allowsKeyedCoding():
323            return long(coder.decodeObjectForKey_(kVALUE))
324        else:
325            return long(coder.decodeObject())
326    decode_dispatch[kOP_LONG] = load_long
327
328    def load_float(coder, setValue):
329        if coder.allowsKeyedCoding():
330            return coder.decodeFloatForKey_(kVALUE)
331        else:
332            raise RuntimeError("Unexpected encoding")
333    decode_dispatch[kOP_FLOAT] = load_float
334
335    def load_float_str(coder, setValue):
336        if coder.allowsKeyedCoding():
337            return float(coder.decodeObjectForKey_(kVALUE))
338        else:
339            return float(coder.decodeObject())
340    decode_dispatch[kOP_FLOAT_STR] = load_float_str
341
342    def load_tuple(coder, setValue):
343        if coder.allowsKeyedCoding():
344            return tuple(coder.decodeObjectForKey_(kVALUE))
345        else:
346            return tuple(coder.decodeObject())
347
348    decode_dispatch[kOP_TUPLE] = load_tuple
349
350    def load_list(coder, setValue):
351        if coder.allowsKeyedCoding():
352            return list(coder.decodeObjectForKey_(kVALUE))
353        else:
354            return list(coder.decodeObject())
355    decode_dispatch[kOP_LIST] = load_list
356
357    def load_dict(coder, setValue):
358        if coder.allowsKeyedCoding():
359            return dict(coder.decodeObjectForKey_(kVALUE))
360        else:
361            return dict(coder.decodeObject())
362    decode_dispatch[kOP_DICT] = load_dict
363
364    def load_global_ext(coder, setValue):
365        if coder.allowsKeyedCoding():
366            code = coder.intForKey_(kCODE)
367        else:
368            code = coder.__pyobjc__decodeInt()
369        nil = []
370        obj = copy_reg._extension_cache.get(code, nil)
371        if obj is not nil:
372            return obj
373        key = copy_reg._inverted_registry.get(code)
374        if not key:
375            raise ValueError("unregistered extension code %d" % code)
376
377        module, name = key
378        __import__(module)
379        mod = sys.modules[module]
380        klass = getattr(mod, name)
381        return klass
382    decode_dispatch[kOP_GLOBAL_EXT] = load_global_ext
383
384    def load_global(coder, setValue):
385        if coder.allowsKeyedCoding():
386            module = coder.decodeObjectForKey_(kMODULE)
387            name = coder.decodeObjectForKey_(kNAME)
388        else:
389            module = coder.decodeObject()
390            name = coder.decodeObject()
391
392        __import__(module)
393        mod = sys.modules[module]
394        klass = getattr(mod, name)
395        return klass
396
397    decode_dispatch[kOP_GLOBAL] = load_global
398
399
400    def load_inst(coder, setValue):
401        if coder.allowsKeyedCoding():
402            cls = coder.decodeObjectForKey_(kCLASS)
403            initargs = coder.decodeObjectForKey_(kARGS)
404        else:
405            cls = coder.decodeObject()
406            initargs = coder.decodeObject()
407
408
409        instantiated = 0
410        if (sys.version_info[0] == 2 and not initargs and
411                type(cls) is ClassType and
412                not hasattr(cls, "__getinitargs__")):
413            try:
414                value = _EmptyClass()
415                value.__class__ = cls
416                instantiated = 1
417
418            except RuntimeError:
419                pass
420
421        if not instantiated:
422            try:
423                value = cls(*initargs)
424            except TypeError, err:
425                raise TypeError, "in constructor for %s: %s" % (
426                    cls.__name__, str(err)), sys.exc_info()[2]
427
428
429        # We now have the object, but haven't set the correct
430        # state yet.  Tell the bridge about this value right
431        # away, that's needed because `value` might be part
432        # of the object state which we'll retrieve next.
433        setValue(value)
434
435        if coder.allowsKeyedCoding():
436            state = coder.decodeObjectForKey_(kSTATE)
437        else:
438            state = coder.decodeObject()
439        setstate = getattr(value, "__setstate__", None)
440        if setstate is not None:
441            setstate(state)
442            return value
443
444        slotstate = None
445        if isinstance(state, tuple) and len(state) == 2:
446            state, slotstate = state
447
448        if state:
449            try:
450                inst_dict = value.__dict__
451                for k, v in state.iteritems():
452                    if type(k) == str:
453                        inst_dict[intern(k)] = v
454                    else:
455                        inst_dict[k] = v
456
457            except RuntimeError:
458                for k, v in state.items():
459                    setattr(value, intern(k), v)
460
461        if slotstate:
462            for k, v in slotstate.items():
463                setattr(value, intern(k), v)
464
465        return value
466    decode_dispatch[kOP_INST] = load_inst
467
468
469    def load_reduce(coder, setValue):
470        if coder.allowsKeyedCoding():
471            func = coder.decodeObjectForKey_(kFUNC)
472            args = coder.decodeObjectForKey_(kARGS)
473        else:
474            func = coder.decodeObject()
475            args = coder.decodeObject()
476
477        value = func(*args)
478
479        # We now have the object, but haven't set the correct
480        # state yet.  Tell the bridge about this value right
481        # away, that's needed because `value` might be part
482        # of the object state which we'll retrieve next.
483        setValue(value)
484
485        if coder.allowsKeyedCoding():
486            listitems = coder.decodeObjectForKey_(kLIST)
487            dictitems = coder.decodeObjectForKey_(kDICT)
488            state = coder.decodeObjectForKey_(kSTATE)
489        else:
490            listitems = coder.decodeObject()
491            dictitems = coder.decodeObject()
492            state = coder.decodeObject()
493
494        setstate = getattr(value, "__setstate__", None)
495        if setstate:
496            setstate(state)
497            return
498
499        slotstate = None
500        if isinstance(state, tuple) and len(state) == 2:
501            state, slotstate = state
502
503        if state:
504            try:
505                inst_dict = value.__dict__
506
507                for k, v in state.iteritems():
508                    if type(k) == str:
509                        inst_dict[intern(k)] = v
510                    else:
511                        inst_dict[k] = v
512
513            except RuntimeError:
514                for k, v in state.items():
515                    setattr(value, intern(k), v)
516
517        if slotstate:
518            for k, v in slotstate.items():
519                setattr(value, intern(k), v)
520
521        if listitems:
522            for a in listitems:
523                value.append(a)
524
525        if dictitems:
526            for k, v in dictitems.items():
527                value[k] = v
528
529        return value
530    decode_dispatch[kOP_REDUCE] = load_reduce
531
532
533    def pyobjectEncode(self, coder):
534        t = type(self)
535
536        # Find builtin support
537        f = encode_dispatch.get(t)
538        if f is not None:
539            f(coder, self)
540            return
541
542        # Check for a class with a custom metaclass
543        try:
544            issc = issubclass(t, type)
545        except TypeError:
546            issc = 0
547
548        if issc:
549            save_global(coder, self)
550            return
551
552        # Check copy_reg.dispatch_table
553        reduce = copy_reg.dispatch_table.get(t)
554        if reduce is not None:
555            rv = reduce(self)
556
557        else:
558            reduce = getattr(self, "__reduce_ex__", None)
559            if reduce is not None:
560                rv = reduce(2)
561
562            else:
563                rv = getattr(self, "__reduce__", None)
564                if reduce is not None:
565                    rv = reduce()
566
567                else:
568                    raise PicklingError("Can't pickle %r object: %r" %
569                            (t.__name__, self))
570
571        if type(rv) is str:
572            save_global(coder, rv)
573            return
574
575        if type(rv) is not tuple:
576            raise PicklingError("%s must return string or tuple" % reduce)
577
578        l = len(rv)
579        if not (2 <= l <= 5):
580            raise PicklingError("Tuple returned by %s must have two to "
581                    "five elements" % reduce)
582
583        save_reduce(coder, *rv)
584
585    def pyobjectDecode(coder, setValue):
586        if coder.allowsKeyedCoding():
587            tp = coder.decodeIntForKey_(kKIND)
588        else:
589            tp = coder.__pyobjc__decodeInt()
590        f = decode_dispatch.get(tp)
591        if f is None:
592            raise UnpicklingError("Unknown object kind: %s"%(tp,))
593
594        return f(coder, setValue)
595
596    # An finally register the coder/decoder
597    OC_PythonObject.setVersion_coder_decoder_copier_(
598            1, pyobjectEncode, pyobjectDecode, copy.copy)
599
600setupPythonObject()
601