Home | History | Annotate | Download | only in test 
      1 import pickle
      2 from cStringIO import StringIO
      3 
      4 from test import test_support
      5 
      6 from test.pickletester import (AbstractPickleTests,
      7                                AbstractPickleModuleTests,
      8                                AbstractPersistentPicklerTests,
      9                                AbstractPicklerUnpicklerObjectTests,
     10                                BigmemPickleTests)
     11 
     12 class PickleTests(AbstractPickleTests, AbstractPickleModuleTests):
     13 
     14     def dumps(self, arg, proto=0, fast=0):
     15         # Ignore fast
     16         return pickle.dumps(arg, proto)
     17 
     18     def loads(self, buf):
     19         # Ignore fast
     20         return pickle.loads(buf)
     21 
     22     module = pickle
     23     error = KeyError
     24 
     25 class PicklerTests(AbstractPickleTests):
     26 
     27     error = KeyError
     28 
     29     def dumps(self, arg, proto=0, fast=0):
     30         f = StringIO()
     31         p = pickle.Pickler(f, proto)
     32         if fast:
     33             p.fast = fast
     34         p.dump(arg)
     35         f.seek(0)
     36         return f.read()
     37 
     38     def loads(self, buf):
     39         f = StringIO(buf)
     40         u = pickle.Unpickler(f)
     41         return u.load()
     42 
     43 class PersPicklerTests(AbstractPersistentPicklerTests):
     44 
     45     def dumps(self, arg, proto=0, fast=0):
     46         class PersPickler(pickle.Pickler):
     47             def persistent_id(subself, obj):
     48                 return self.persistent_id(obj)
     49         f = StringIO()
     50         p = PersPickler(f, proto)
     51         if fast:
     52             p.fast = fast
     53         p.dump(arg)
     54         f.seek(0)
     55         return f.read()
     56 
     57     def loads(self, buf):
     58         class PersUnpickler(pickle.Unpickler):
     59             def persistent_load(subself, obj):
     60                 return self.persistent_load(obj)
     61         f = StringIO(buf)
     62         u = PersUnpickler(f)
     63         return u.load()
     64 
     65 class PicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
     66 
     67     pickler_class = pickle.Pickler
     68     unpickler_class = pickle.Unpickler
     69 
     70 class PickleBigmemPickleTests(BigmemPickleTests):
     71 
     72     def dumps(self, arg, proto=0, fast=0):
     73         # Ignore fast
     74         return pickle.dumps(arg, proto)
     75 
     76     def loads(self, buf):
     77         # Ignore fast
     78         return pickle.loads(buf)
     79 
     80 
     81 def test_main():
     82     test_support.run_unittest(
     83         PickleTests,
     84         PicklerTests,
     85         PersPicklerTests,
     86         PicklerUnpicklerObjectTests,
     87         PickleBigmemPickleTests,
     88     )
     89     test_support.run_doctest(pickle)
     90 
     91 if __name__ == "__main__":
     92     test_main()
     93