Lib/pickle.py

   1 """\
   2 Pickling Algorithm
   3 ------------------
   4
   5 This module implements a basic but powerful algorithm for "pickling" (a.k.a.
   6 serializing, marshalling or flattening) nearly arbitrary Python objects.
   7 This is a more primitive notion than persistency -- although pickle
   8 reads and writes file objects, it does not handle the issue of naming
   9 persistent objects, nor the (even more complicated) area of concurrent
  10 access to persistent objects.  The pickle module can transform a complex
  11 object into a byte stream and it can transform the byte stream into
  12 an object with the same internal structure.  The most obvious thing to
  13 do with these byte streams is to write them onto a file, but it is also
  14 conceivable to send them across a network or store them in a database.
  15
  16 Unlike the built-in marshal module, pickle handles the following correctly:
  17
  18 - recursive objects
  19 - pointer sharing
  20 - classes and class instances
  21
  22 Pickle is Python-specific.  This has the advantage that there are no
  23 restrictions imposed by external standards such as CORBA (which probably
  24 can't represent pointer sharing or recursive objects); however it means
  25 that non-Python programs may not be able to reconstruct pickled Python
  26 objects.
  27
  28 Pickle uses a printable ASCII representation.  This is slightly more
  29 voluminous than a binary representation.  However, small integers actually
  30 take *less* space when represented as minimal-size decimal strings than
  31 when represented as 32-bit binary numbers, and strings are only much longer
  32 if they contain control characters or 8-bit characters.  The big advantage
  33 of using printable ASCII (and of some other characteristics of pickle's
  34 representation) is that for debugging or recovery purposes it is possible
  35 for a human to read the pickled file with a standard text editor.  (I could
  36 have gone a step further and used a notation like S-expressions, but the
  37 parser would have been considerably more complicated and slower, and the
  38 files would probably have become much larger.)
  39
  40 Pickle doesn't handle code objects, which marshal does.
  41 I suppose pickle could, and maybe it should, but there's probably no
  42 great need for it right now (as long as marshal continues to be used
  43 for reading and writing code objects), and at least this avoids
  44 the possibility of smuggling Trojan horses into a program.
  45
  46 For the benefit of persistency modules written using pickle, it supports
  47 the notion of a reference to an object outside the pickled data stream.
  48 Such objects are referenced by a name, which is an arbitrary string of
  49 printable ASCII characters.  The resolution of such names is not defined
  50 by the pickle module -- the persistent object module will have to implement
  51 a method "persistent_load".  To write references to persistent objects,
  52 the persistent module must define a method "persistent_id" which returns
  53 either None or the persistent ID of the object.
  54
  55 There are some restrictions on the pickling of class instances.
  56
  57 First of all, the class must be defined at the top level in a module.
  58
  59 Next, it must normally be possible to create class instances by calling
  60 the class without arguments.  If this is undesirable, the class can
  61 define a method __getinitargs__ (XXX not a pretty name!), which should
  62 return a *tuple* containing the arguments to be passed to the class
  63 constructor.
  64
  65 Classes can influence how their instances are pickled -- if the class defines
  66 the method __getstate__, it is called and the return state is pickled
  67 as the contents for the instance, and if the class defines the
  68 method __setstate__, it is called with the unpickled state.  (Note
  69 that these methods can also be used to implement copying class instances.)
  70 If there is no __getstate__ method, the instance's __dict__
  71 is pickled.  If there is no __setstate__ method, the pickled object
  72 must be a dictionary and its items are assigned to the new instance's
  73 dictionary.  (If a class defines both __getstate__ and __setstate__,
  74 the state object needn't be a dictionary -- these methods can do what they
  75 want.)
  76
  77 Note that when class instances are pickled, their class's code and data
  78 is not pickled along with them.  Only the instance data is pickled.
  79 This is done on purpose, so you can fix bugs in a class or add methods and
  80 still load objects that were created with an earlier version of the
  81 class.  If you plan to have long-lived objects that will see many versions
  82 of a class, it may be worth to put a version number in the objects so
  83 that suitable conversions can be made by the class's __setstate__ method.
  84
  85 The interface is as follows:
  86
  87 To pickle an object x onto a file f, open for writing:
  88
  89         p = pickle.Pickler(f)
  90         p.dump(x)
  91
  92 To unpickle an object x from a file f, open for reading:
  93
  94         u = pickle.Unpickler(f)
  95         x = u.load()
  96
  97 The Pickler class only calls the method f.write with a string argument
  98 (XXX possibly the interface should pass f.write instead of f).
  99 The Unpickler calls the methods f.read(with an integer argument)
 100 and f.readline(without argument), both returning a string.
 101 It is explicitly allowed to pass non-file objects here, as long as they
 102 have the right methods.
 103
 104 The following types can be pickled:
 105
 106 - None
 107 - integers, long integers, floating point numbers
 108 - strings
 109 - tuples, lists and dictionaries containing only picklable objects
 110 - class instances whose __dict__ or __setstate__() is picklable
 111 - classes
 112
 113 Attempts to pickle unpicklable objects will raise an exception
 114 after having written an unspecified number of bytes to the file argument.
 115
 116 It is possible to make multiple calls to Pickler.dump() or to
 117 Unpickler.load(), as long as there is a one-to-one correspondence
 118 between pickler and Unpickler objects and between dump and load calls
 119 for any pair of corresponding Pickler and Unpicklers.  WARNING: this
 120 is intended for pickleing multiple objects without intervening modifications
 121 to the objects or their parts.  If you modify an object and then pickle
 122 it again using the same Pickler instance, the object is not pickled
 123 again -- a reference to it is pickled and the Unpickler will return
 124 the old value, not the modified one.  (XXX There are two problems here:
 125 (a) detecting changes, and (b) marshalling a minimal set of changes.
 126 I have no answers.  Garbage Collection may also become a problem here.)
 127 """
 128
 129 __version__ = "1.5"                     # Code version
 130
 131 from types import *
 132 import string
 133
 134 format_version = "1.1"                  # File format version we write
 135 compatible_formats = ["1.0"]            # Old format versions we can read
 136
 137 PicklingError = "pickle.PicklingError"
 138
 139 AtomicTypes = [NoneType, IntType, FloatType, StringType]
 140
 141 def safe(object):
 142         t = type(object)
 143         if t in AtomicTypes:
 144                 return 1
 145         if t is TupleType:
 146                 for item in object:
 147                         if not safe(item): return 0
 148                 return 1
 149         return 0
 150
 151 MARK = '('
 152 POP = '0'
 153 DUP = '2'
 154 STOP = '.'
 155 TUPLE = 't'
 156 LIST = 'l'
 157 DICT = 'd'
 158 INST = 'i'
 159 CLASS = 'c'
 160 GET = 'g'
 161 PUT = 'p'
 162 APPEND = 'a'
 163 SETITEM = 's'
 164 BUILD = 'b'
 165 NONE = 'N'
 166 INT = 'I'
 167 LONG = 'L'
 168 FLOAT = 'F'
 169 STRING = 'S'
 170 PERSID = 'P'
 171 AtomicKeys = [NONE, INT, LONG, FLOAT, STRING]
 172 AtomicMap = {
 173         NoneType: NONE,
 174         IntType: INT,
 175         LongType: LONG,
 176         FloatType: FLOAT,
 177         StringType: STRING,
 178 }
 179
 180 class Pickler:
 181
 182         def __init__(self, file):
 183                 self.write = file.write
 184                 self.memo = {}
 185
 186         def dump(self, object):
 187                 self.save(object)
 188                 self.write(STOP)
 189
 190         def save(self, object):
 191                 pid = self.persistent_id(object)
 192                 if pid:
 193                         self.write(PERSID + str(pid) + '\n')
 194                         return
 195                 d = id(object)
 196                 if self.memo.has_key(d):
 197                         self.write(GET + `d` + '\n')
 198                         return
 199                 t = type(object)
 200                 try:
 201                         f = self.dispatch[t]
 202                 except KeyError:
 203                         raise PicklingError, \
 204                               "can't pickle %s objects" % `t.__name__`
 205                 f(self, object)
 206
 207         def persistent_id(self, object):
 208                 return None
 209
 210         dispatch = {}
 211
 212         def save_none(self, object):
 213                 self.write(NONE)
 214         dispatch[NoneType] = save_none
 215
 216         def save_int(self, object):
 217                 self.write(INT + `object` + '\n')
 218         dispatch[IntType] = save_int
 219
 220         def save_long(self, object):
 221                 self.write(LONG + `object` + '\n')
 222         dispatch[LongType] = save_long
 223
 224         def save_float(self, object):
 225                 self.write(FLOAT + `object` + '\n')
 226         dispatch[FloatType] = save_float
 227
 228         def save_string(self, object):
 229                 d = id(object)
 230                 self.write(STRING + `object` + '\n')
 231                 self.write(PUT + `d` + '\n')
 232                 self.memo[d] = object
 233         dispatch[StringType] = save_string
 234
 235         def save_tuple(self, object):
 236                 d = id(object)
 237                 self.write(MARK)
 238                 n = len(object)
 239                 for k in range(n):
 240                         self.save(object[k])
 241                         if self.memo.has_key(d):
 242                                 # Saving object[k] has saved us!
 243                                 while k >= 0:
 244                                         self.write(POP)
 245                                         k = k-1
 246                                 self.write(GET + `d` + '\n')
 247                                 break
 248                 else:
 249                         self.write(TUPLE + PUT + `d` + '\n')
 250                         self.memo[d] = object
 251         dispatch[TupleType] = save_tuple
 252
 253         def save_list(self, object):
 254                 d = id(object)
 255                 self.write(MARK)
 256                 n = len(object)
 257                 for k in range(n):
 258                         item = object[k]
 259                         if not safe(item):
 260                                 break
 261                         self.save(item)
 262                 else:
 263                         k = n
 264                 self.write(LIST + PUT + `d` + '\n')
 265                 self.memo[d] = object
 266                 for k in range(k, n):
 267                         item = object[k]
 268                         self.save(item)
 269                         self.write(APPEND)
 270         dispatch[ListType] = save_list
 271
 272         def save_dict(self, object):
 273                 d = id(object)
 274                 self.write(MARK)
 275                 items = object.items()
 276                 n = len(items)
 277                 for k in range(n):
 278                         key, value = items[k]
 279                         if not safe(key) or not safe(value):
 280                                 break
 281                         self.save(key)
 282                         self.save(value)
 283                 else:
 284                         k = n
 285                 self.write(DICT + PUT + `d` + '\n')
 286                 self.memo[d] = object
 287                 for k in range(k, n):
 288                         key, value = items[k]
 289                         self.save(key)
 290                         self.save(value)
 291                         self.write(SETITEM)
 292         dispatch[DictionaryType] = save_dict
 293
 294         def save_inst(self, object):
 295                 d = id(object)
 296                 cls = object.__class__
 297                 module = whichmodule(cls)
 298                 name = cls.__name__
 299                 if hasattr(object, '__getinitargs__'):
 300                         args = object.__getinitargs__()
 301                         len(args) # XXX Assert it's a sequence
 302                 else:
 303                         args = ()
 304                 self.write(MARK)
 305                 for arg in args:
 306                         self.save(arg)
 307                 self.write(INST + module + '\n' + name + '\n' +
 308                            PUT + `d` + '\n')
 309                 self.memo[d] = object
 310                 try:
 311                         getstate = object.__getstate__
 312                 except AttributeError:
 313                         stuff = object.__dict__
 314                 else:
 315                         stuff = getstate()
 316                 self.save(stuff)
 317                 self.write(BUILD)
 318         dispatch[InstanceType] = save_inst
 319
 320         def save_class(self, object):
 321                 d = id(object)
 322                 module = whichmodule(object)
 323                 name = object.__name__
 324                 self.write(CLASS + module + '\n' + name + '\n' +
 325                            PUT + `d` + '\n')
 326         dispatch[ClassType] = save_class
 327
 328
 329 classmap = {}
 330
 331 def whichmodule(cls):
 332         """Figure out the module in which a class occurs.
 333
 334         Search sys.modules for the module.
 335         Cache in classmap.
 336         Return a module name.
 337         If the class cannot be found, return __main__.
 338         """
 339         if classmap.has_key(cls):
 340                 return classmap[cls]
 341         import sys
 342         clsname = cls.__name__
 343         for name, module in sys.modules.items():
 344                 if name != '__main__' and \
 345                    hasattr(module, clsname) and \
 346                    getattr(module, clsname) is cls:
 347                         break
 348         else:
 349                 name = '__main__'
 350         classmap[cls] = name
 351         return name
 352
 353
 354 class Unpickler:
 355
 356         def __init__(self, file):
 357                 self.readline = file.readline
 358                 self.read = file.read
 359                 self.memo = {}
 360
 361         def load(self):
 362                 self.mark = ['spam'] # Any new unique object
 363                 self.stack = []
 364                 try:
 365                         while 1:
 366                                 key = self.read(1)
 367                                 self.dispatch[key](self)
 368                 except STOP, value:
 369                         return value
 370
 371         def marker(self):
 372                 k = len(self.stack)-1
 373                 while self.stack[k] != self.mark: k = k-1
 374                 return k
 375
 376         dispatch = {}
 377
 378         def load_eof(self):
 379                 raise EOFError
 380         dispatch[''] = load_eof
 381
 382         def load_persid(self):
 383                 pid = self.readline()[:-1]
 384                 self.stack.append(self.persistent_load(pid))
 385         dispatch[PERSID] = load_persid
 386
 387         def load_none(self):
 388                 self.stack.append(None)
 389         dispatch[NONE] = load_none
 390
 391         def load_atomic(self):
 392                 self.stack.append(eval(self.readline()[:-1]))
 393         dispatch[INT] = load_atomic
 394         dispatch[LONG] = load_atomic
 395         dispatch[FLOAT] = load_atomic
 396         dispatch[STRING] = load_atomic
 397
 398         def load_tuple(self):
 399                 k = self.marker()
 400                 self.stack[k:] = [tuple(self.stack[k+1:])]
 401         dispatch[TUPLE] = load_tuple
 402
 403         def load_list(self):
 404                 k = self.marker()
 405                 self.stack[k:] = [self.stack[k+1:]]
 406         dispatch[LIST] = load_list
 407
 408         def load_dict(self):
 409                 k = self.marker()
 410                 d = {}
 411                 items = self.stack[k+1:]
 412                 for i in range(0, len(items), 2):
 413                         key = items[i]
 414                         value = items[i+1]
 415                         d[key] = value
 416                 self.stack[k:] = [d]
 417         dispatch[DICT] = load_dict
 418
 419         def load_inst(self):
 420                 k = self.marker()
 421                 args = tuple(self.stack[k+1:])
 422                 del self.stack[k:]
 423                 module = self.readline()[:-1]
 424                 name = self.readline()[:-1]
 425                 klass = self.find_class(module, name)
 426                 value = apply(klass, args)
 427                 self.stack.append(value)
 428         dispatch[INST] = load_inst
 429
 430         def load_class(self):
 431                 module = self.readline()[:-1]
 432                 name = self.readline()[:-1]
 433                 klass = self.find_class(module, name)
 434                 self.stack.append(klass)
 435                 return klass
 436         dispatch[CLASS] = load_class
 437
 438         def find_class(self, module, name):
 439                 env = {}
 440                 try:
 441                         exec 'from %s import %s' % (module, name) in env
 442                 except ImportError:
 443                         raise SystemError, \
 444                               "Failed to import class %s from module %s" % \
 445                               (name, module)
 446                 klass = env[name]
 447                 if type(klass) != ClassType:
 448                         raise SystemError, \
 449                          "Imported object %s from module %s is not a class" % \
 450                          (name, module)
 451                 return klass
 452
 453         def load_pop(self):
 454                 del self.stack[-1]
 455         dispatch[POP] = load_pop
 456
 457         def load_dup(self):
 458                 stack.append(stack[-1])
 459         dispatch[DUP] = load_dup
 460
 461         def load_get(self):
 462                 self.stack.append(self.memo[string.atoi(self.readline()[:-1])])
 463         dispatch[GET] = load_get
 464
 465         def load_put(self):
 466                 self.memo[string.atoi(self.readline()[:-1])] = self.stack[-1]
 467         dispatch[PUT] = load_put
 468
 469         def load_append(self):
 470                 value = self.stack[-1]
 471                 del self.stack[-1]
 472                 list = self.stack[-1]
 473                 list.append(value)
 474         dispatch[APPEND] = load_append
 475
 476         def load_setitem(self):
 477                 value = self.stack[-1]
 478                 key = self.stack[-2]
 479                 del self.stack[-2:]
 480                 dict = self.stack[-1]
 481                 dict[key] = value
 482         dispatch[SETITEM] = load_setitem
 483
 484         def load_build(self):
 485                 value = self.stack[-1]
 486                 del self.stack[-1]
 487                 inst = self.stack[-1]
 488                 try:
 489                         setstate = inst.__setstate__
 490                 except AttributeError:
 491                         for key in value.keys():
 492                                 inst.__dict__[key] = value[key]
 493                 else:
 494                         setstate(value)
 495         dispatch[BUILD] = load_build
 496
 497         def load_mark(self):
 498                 self.stack.append(self.mark)
 499         dispatch[MARK] = load_mark
 500
 501         def load_stop(self):
 502                 value = self.stack[-1]
 503                 del self.stack[-1]
 504                 raise STOP, value
 505         dispatch[STOP] = load_stop
 506
 507
 508 # Shorthands
 509
 510 def dump(object, file):
 511         Pickler(file).dump(object)
 512
 513 def dumps(object):
 514         import StringIO
 515         file = StringIO.StringIO()
 516         Pickler(file).dump(object)
 517         return file.getvalue()
 518
 519 def load(file):
 520         return Unpickler(file).load()
 521
 522 def loads(str):
 523         import StringIO
 524         file = StringIO.StringIO(str)
 525         return Unpickler(file).load()
 526
 527
 528 # The rest is used for testing only
 529
 530 class C:
 531         def __cmp__(self, other):
 532                 return cmp(self.__dict__, other.__dict__)
 533
 534 def test():
 535         fn = 'pickle_tmp'
 536         c = C()
 537         c.foo = 1
 538         c.bar = 2
 539         x = [0,1,2,3]
 540         y = ('abc', 'abc', c, c)
 541         x.append(y)
 542         x.append(y)
 543         x.append(5)
 544         f = open(fn, 'w')
 545         F = Pickler(f)
 546         F.dump(x)
 547         f.close()
 548         f = open(fn, 'r')
 549         U = Unpickler(f)
 550         x2 = U.load()
 551         print x
 552         print x2
 553         print x == x2
 554         print map(id, x)
 555         print map(id, x2)
 556         print F.memo
 557         print U.memo
 558
 559 if __name__ == '__main__':
 560         test()