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