diff --git a/PickleCompress/PickleCompression.py b/PickleCompress/PickleCompression.py new file mode 100644 index 0000000..66317de --- /dev/null +++ b/PickleCompress/PickleCompression.py @@ -0,0 +1,73 @@ +import bz2 +import lzma +import gzip +import pickle + + +class CompressPickle: + def compress_pickle(self, filename: str, object, compression_mode: str = "none"): + self.object = object + if "." in filename: + filename = filename.split(".")[0] + if compression_mode == "gzip": + filename = f"{filename}.gz" + with gzip.open(filename=filename, mode="wb") as gzfile: + pickle.dump(obj=self.object, file=gzfile) + elif compression_mode == "bzip2": + filename = f"{filename}.bz2" + with bz2.open(filename=filename, mode="wb") as bz2file: + pickle.dump(obj=self.object, file=bz2file) + elif compression_mode == "lzma": + filename = f"{filename}.xz" + with lzma.open(filename=filename, mode="wb") as lzfile: + pickle.dump(obj=self.object, file=lzfile) + elif compression_mode == "none": + filename = f"{filename}.pickle" + with open(filename, mode="wb") as picklefile: + pickle.dump(obj=self.object, file=picklefile) + else: + print("Incorrect Format") + + + def decompress_pickle(self,filename: str): + if "gz" in filename: + with gzip.open(filename=filename, mode="rb") as gzfile: + self.object = pickle.load(file=gzfile) + return self.object + elif "bz2" in filename: + with bz2.open(filename=filename, mode="rb") as bz2file: + self.object = pickle.load(file=bz2file) + return self.object + elif "xz" in filename: + with lzma.open(filename=filename, mode="rb") as lzfile: + self.object = pickle.load(file=lzfile) + return self.object + elif "pickle" in filename: + with open(filename, mode="rb") as picklefile: + self.object = pickle.load(file=picklefile) + return self.object + else: + print("Incorrect Format") + + +# Main Code Here +def main(): + a: str = "Test" + + compressor = CompressPickle() + # Compress to each mode 'gzip', 'bzip2', and 'lzma' + compressor.compress_pickle(filename="a", object=a, compression_mode="gzip") + compressor.compress_pickle(filename="a", object=a, compression_mode="bzip2") + compressor.compress_pickle(filename="a", object=a, compression_mode="lzma") + compressor.compress_pickle(filename="a", object=a) + + # Decompress + b = compressor.decompress_pickle(filename="a.gz") + c = compressor.decompress_pickle(filename="a.bz2") + d = compressor.decompress_pickle(filename="a.xz") + e = compressor.decompress_pickle(filename="a.pickle") + + print(a, b, c, d, e) + +if __name__ == "__main__": + main() diff --git a/README.md b/README.md index 5381800..4c91990 100644 --- a/README.md +++ b/README.md @@ -3,5 +3,6 @@ This repository contains useful programs written in **Python 3.9**. Each folder accompanied with **requirements.txt** for easier dependency access. ## Currently Avaiable: ## -* Benchmark Programs : using **psutil** to obtain performance +* Benchmark Program : using **psutil** to obtain performance +* Pickle Compress Program : compress variable dumps using **gzip**, **bzip2**, **lzma**, or **no compression at all** * *adding more soon*