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