pickle.py

mallet是自然语言处理、机器学习领域的一个开源项目。

"""Create portable serialized representations of Python objects.

See module cPickle for a (much) faster implementation.
See module copy_reg for a mechanism for registering custom picklers.

Classes:

    Pickler
    Unpickler

Functions:

    dump(object, file)
    dumps(object) -> string
    load(file) -> object
    loads(string) -> object

Misc variables:

    __version__
    format_version
    compatible_formats

"""

__version__ = "$Revision: 1.1 $"       # Code version

from types import *
from copy_reg import dispatch_table, safe_constructors
import marshal
import sys
import struct
import re

__all__ = ["PickleError", "PicklingError", "UnpicklingError", "Pickler",
           "Unpickler", "dump", "dumps", "load", "loads"]

format_version = "1.3"                     # File format version we write
compatible_formats = ["1.0", "1.1", "1.2"] # Old format versions we can read

mdumps = marshal.dumps
mloads = marshal.loads

class PickleError(Exception): pass
class PicklingError(PickleError): pass
class UnpicklingError(PickleError): pass

class _Stop(Exception):
    def __init__(self, value):
        self.value = value

try:
    from org.python.core import PyStringMap
except ImportError:
    PyStringMap = None

try:
    UnicodeType
except NameError:
    UnicodeType = None


MARK            = '('
STOP            = '.'
POP             = '0'
POP_MARK        = '1'
DUP             = '2'
FLOAT           = 'F'
INT             = 'I'
BININT          = 'J'
BININT1         = 'K'
LONG            = 'L'
BININT2         = 'M'
NONE            = 'N'
PERSID          = 'P'
BINPERSID       = 'Q'
REDUCE          = 'R'
STRING          = 'S'
BINSTRING       = 'T'
SHORT_BINSTRING = 'U'
UNICODE         = 'V'
BINUNICODE      = 'X'
APPEND          = 'a'
BUILD           = 'b'
GLOBAL          = 'c'
DICT            = 'd'
EMPTY_DICT      = '}'
APPENDS         = 'e'
GET             = 'g'
BINGET          = 'h'
INST            = 'i'
LONG_BINGET     = 'j'
LIST            = 'l'
EMPTY_LIST      = ']'
OBJ             = 'o'
PUT             = 'p'
BINPUT          = 'q'
LONG_BINPUT     = 'r'
SETITEM         = 's'
TUPLE           = 't'
EMPTY_TUPLE     = ')'
SETITEMS        = 'u'
BINFLOAT        = 'G'

__all__.extend([x for x in dir() if re.match("[A-Z][A-Z0-9_]+$",x)])

class Pickler:

    def __init__(self, file, bin = 0):
        self.write = file.write
        self.memo = {}
        self.bin = bin

    def dump(self, object):
        self.save(object)
        self.write(STOP)

    def put(self, i):
        if self.bin:
            s = mdumps(i)[1:]
            if i < 256:
                return BINPUT + s[0]

            return LONG_BINPUT + s

        return PUT + `i` + '\n'

    def get(self, i):
        if self.bin:
            s = mdumps(i)[1:]

            if i < 256:
                return BINGET + s[0]

            return LONG_BINGET + s

        return GET + `i` + '\n'

    def save(self, object, pers_save = 0):
        memo = self.memo

        if not pers_save:
            pid = self.persistent_id(object)
            if pid is not None:
                self.save_pers(pid)
                return

        d = id(object)

        t = type(object)

        if (t is TupleType) and (len(object) == 0):
            if self.bin:
                self.save_empty_tuple(object)
            else:
                self.save_tuple(object)
            return

        if memo.has_key(d):
            self.write(self.get(memo[d][0]))
            return

        try:
            f = self.dispatch[t]
        except KeyError:
            pid = self.inst_persistent_id(object)
            if pid is not None:
                self.save_pers(pid)
                return

            try:
                issc = issubclass(t, TypeType)
            except TypeError: # t is not a class
                issc = 0
            if issc:
                self.save_global(object)
                return

            try:
                reduce = dispatch_table[t]
            except KeyError:
                try:
                    reduce = object.__reduce__
                except AttributeError:
                    raise PicklingError, \
                        "can't pickle %s object: %s" % (`t.__name__`,
                                                         `object`)
                else:
                    tup = reduce()
            else:
                tup = reduce(object)

            if type(tup) is StringType:
                self.save_global(object, tup)
                return

            if type(tup) is not TupleType:
                raise PicklingError, "Value returned by %s must be a " \
                                     "tuple" % reduce

            l = len(tup)

            if (l != 2) and (l != 3):
                raise PicklingError, "tuple returned by %s must contain " \
                                     "only two or three elements" % reduce

            callable = tup[0]
            arg_tup  = tup[1]

            if l > 2:
                state = tup[2]
            else:
                state = None

            if type(arg_tup) is not TupleType and arg_tup is not None:
                raise PicklingError, "Second element of tuple returned " \
                                     "by %s must be a tuple" % reduce

            self.save_reduce(callable, arg_tup, state)
            memo_len = len(memo)
            self.write(self.put(memo_len))
            memo[d] = (memo_len, object)
            return

        f(self, object)

    def persistent_id(self, object):
        return None

    def inst_persistent_id(self, object):
        return None

    def save_pers(self, pid):
        if not self.bin:
            self.write(PERSID + str(pid) + '\n')
        else:
            self.save(pid, 1)
            self.write(BINPERSID)

    def save_reduce(self, callable, arg_tup, state = None):
        write = self.write
        save = self.save

        save(callable)
        save(arg_tup)
        write(REDUCE)

        if state is not None:
            save(state)
            write(BUILD)

    dispatch = {}

    def save_none(self, object):
        self.write(NONE)
    dispatch[NoneType] = save_none

    def save_int(self, object):
        if self.bin:
            # If the int is small enough to fit in a signed 4-byte 2's-comp
            # format, we can store it more efficiently than the general
            # case.
            high_bits = object >> 31  # note that Python shift sign-extends
            if  high_bits == 0 or high_bits == -1:
                # All high bits are copies of bit 2**31, so the value
                # fits in a 4-byte signed int.
                i = mdumps(object)[1:]
                assert len(i) == 4
                if i[-2:] == '\000\000':    # fits in 2-byte unsigned int
                    if i[-3] == '\000':     # fits in 1-byte unsigned int
                        self.write(BININT1 + i[0])
                    else:
                        self.write(BININT2 + i[:2])
                else:
                    self.write(BININT + i)
                return
        # Text pickle, or int too big to fit in signed 4-byte format.
        self.write(INT + `object` + '\n')
    dispatch[IntType] = save_int

    def save_long(self, object):
        self.write(LONG + `object` + '\n')
    dispatch[LongType] = save_long

    def save_float(self, object, pack=struct.pack):
        if self.bin:
            self.write(BINFLOAT + pack('>d', object))
        else:
            self.write(FLOAT + `object` + '\n')
    dispatch[FloatType] = save_float

    def save_string(self, object):
        d = id(object)
        memo = self.memo

        if self.bin:
            l = len(object)
            s = mdumps(l)[1:]
            if l < 256:
                self.write(SHORT_BINSTRING + s[0] + object)
            else:
                self.write(BINSTRING + s + object)
        else:
            self.write(STRING + `object` + '\n')

        memo_len = len(memo)
        self.write(self.put(memo_len))
        memo[d] = (memo_len, object)
    dispatch[StringType] = save_string

    def save_unicode(self, object):
        d = id(object)
        memo = self.memo

        if self.bin:
            encoding = object.encode('utf-8')
            l = len(encoding)
            s = mdumps(l)[1:]
            self.write(BINUNICODE + s + encoding)
        else:
            object = object.replace("\\", "\\u005c")
            object = object.replace("\n", "\\u000a")
            self.write(UNICODE + object.encode('raw-unicode-escape') + '\n')

        memo_len = len(memo)
        self.write(self.put(memo_len))
        memo[d] = (memo_len, object)
    dispatch[UnicodeType] = save_unicode

    if StringType == UnicodeType:
        # This is true for Jython
        def save_string(self, object):
            d = id(object)
            memo = self.memo
            unicode = object.isunicode()

            if self.bin:
                if unicode:
                    object = object.encode("utf-8")
                l = len(object)
                s = mdumps(l)[1:]
                if l < 256 and not unicode:
                    self.write(SHORT_BINSTRING + s[0] + object)
                else:
                    if unicode:
                        self.write(BINUNICODE + s + object)
                    else:
                        self.write(BINSTRING + s + object)
            else:
                if unicode:
                    object = object.replace("\\", "\\u005c")
                    object = object.replace("\n", "\\u000a")
                    object = object.encode('raw-unicode-escape')
                    self.write(UNICODE + object + '\n')
                else:
                    self.write(STRING + `object` + '\n')

            memo_len = len(memo)
            self.write(self.put(memo_len))
            memo[d] = (memo_len, object)
        dispatch[StringType] = save_string

    def save_tuple(self, object):

        write = self.write
        save  = self.save
        memo  = self.memo

        d = id(object)

        write(MARK)

        for element in object:
            save(element)

        if len(object) and memo.has_key(d):
            if self.bin:
                write(POP_MARK + self.get(memo[d][0]))
                return

            write(POP * (len(object) + 1) + self.get(memo[d][0]))
            return

        memo_len = len(memo)
        self.write(TUPLE + self.put(memo_len))
        memo[d] = (memo_len, object)
    dispatch[TupleType] = save_tuple

    def save_empty_tuple(self, object):
        self.write(EMPTY_TUPLE)

    def save_list(self, object):
        d = id(object)

        write = self.write
        save  = self.save
        memo  = self.memo

        if self.bin:
            write(EMPTY_LIST)
        else:
            write(MARK + LIST)

        memo_len = len(memo)
        write(self.put(memo_len))
        memo[d] = (memo_len, object)

        using_appends = (self.bin and (len(object) > 1))

        if using_appends:
            write(MARK)

        for element in object:
            save(element)

            if not using_appends:
                write(APPEND)

        if using_appends:
            write(APPENDS)
    dispatch[ListType] = save_list

    def save_dict(self, object):
        d = id(object)

        write = self.write
        save  = self.save
        memo  = self.memo

        if self.bin:
            write(EMPTY_DICT)
        else:
            write(MARK + DICT)

        memo_len = len(memo)
        self.write(self.put(memo_len))
        memo[d] = (memo_len, object)

        using_setitems = (self.bin and (len(object) > 1))

        if using_setitems:
            write(MARK)

        items = object.items()
        for key, value in items:
            save(key)
            save(value)

            if not using_setitems:
                write(SETITEM)

        if using_setitems:
            write(SETITEMS)

    dispatch[DictionaryType] = save_dict
    if not PyStringMap is None:
        dispatch[PyStringMap] = save_dict

    def save_inst(self, object):
        d = id(object)
        cls = object.__class__

        memo  = self.memo
        write = self.write
        save  = self.save

        if hasattr(object, '__getinitargs__'):
            args = object.__getinitargs__()
            len(args) # XXX Assert it's a sequence
            _keep_alive(args, memo)
        else:
            args = ()

        write(MARK)

        if self.bin:
            save(cls)

        for arg in args:
            save(arg)

        memo_len = len(memo)
        if self.bin:
            write(OBJ + self.put(memo_len))
        else:
            write(INST + cls.__module__ + '\n' + cls.__name__ + '\n' +
                self.put(memo_len))

        memo[d] = (memo_len, object)

        try:
            getstate = object.__getstate__
        except AttributeError:
            stuff = object.__dict__