Somewhat safer unpickler, only builtins (/usr/lib/python2.7/pickle.py)

somewhat_safer_unpickler.py

# #######################################################################
# Somewhat safer unpickler, only builtins (/usr/lib/python2.7/pickle.py)
# #######################################################################

class LimitedUnpickler(pickle.Unpickler):
    """
    This class inherits all the implementation from the builtin
    `pickle.Unpickler` but modifies the `self.dispatch` dictionary
    used to unpack the serialized structures.
    
    This is not checked to be safe, but in case you are using pickled
    files in production and you are searching for some safer way to load
    them to convert the data to a different format, this class can be handy.
    """
    def __init__(self, file):
        pickle.Unpickler.__init__(self, file)
        safe =\
        { pickle.MARK
        , pickle.STOP
        , pickle.POP
        , pickle.POP_MARK
        , pickle.DUP
        , pickle.FLOAT
        , pickle.INT
        , pickle.BININT
        , pickle.BININT1
        , pickle.LONG
        , pickle.BININT2
        , pickle.NONE
        , pickle.PERSID
        , pickle.BINPERSID
        #, pickle.REDUCE
        , pickle.STRING
        , pickle.BINSTRING
        , pickle.SHORT_BINSTRING
        , pickle.UNICODE
        , pickle.BINUNICODE
        , pickle.APPEND
        #, pickle.BUILD
        #, pickle.GLOBAL
        , pickle.DICT
        , pickle.EMPTY_DICT
        , pickle.APPENDS
        , pickle.GET
        , pickle.BINGET
        , pickle.INST
        , pickle.LONG_BINGET
        , pickle.LIST
        , pickle.EMPTY_LIST
        #, pickle.OBJ
        , pickle.PUT
        , pickle.BINPUT
        , pickle.LONG_BINPUT
        , pickle.SETITEM
        , pickle.TUPLE
        , pickle.EMPTY_TUPLE
        , pickle.SETITEMS
        , pickle.BINFLOAT
        , pickle.TRUE
        , pickle.FALSE
        , pickle.PROTO
        #, pickle.NEWOBJ
        #, pickle.EXT1
        #, pickle.EXT2
        #, pickle.EXT4
        , pickle.TUPLE1
        , pickle.TUPLE2
        , pickle.TUPLE3
        , pickle.NEWTRUE
        , pickle.NEWFALSE
        , pickle.LONG1
        , pickle.LONG4
        , ''  # EOF hardcoded as `dispatch[''] = load_eof` in pickle.py
        }

        self.dispatch = {k: v for k, v in self.dispatch.items() if k in safe}

def pickle_limited_load(file):
    return LimitedUnpickler(file).load()