import bz2
import pickle
import _pickle as cPickle
import sys


class SafeUnpickler(cPickle.Unpickler):
    def find_class(self, module, name):
        # Only allow safe built-in types
        safe_builtins = {
            "builtins": {"dict", "list", "str", "int", "float", "set", "tuple"}
        }
        if module in safe_builtins and name in safe_builtins[module]:
            return super().find_class(module, name)
        raise pickle.UnpicklingError(f"Unsafe module: {module}.{name}")


def safe_load(file_obj):
    return SafeUnpickler(file_obj).load()


# Load any compressed pickle file
def decompress_pickle(filename):
    sys.stderr.write("loading pickle %s...\n" % filename)
    with bz2.BZ2File(filename, "rb") as f:
        return safe_load(f)  # from SafeUnpickler above
