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