1""" 2 salt.utils.gzip 3 ~~~~~~~~~~~~~~~ 4 Helper module for handling gzip consistently between 2.7+ and 2.6- 5""" 6 7 8import gzip 9import io 10 11import salt.utils.files 12 13 14class GzipFile(gzip.GzipFile): 15 def __init__(self, filename=None, mode=None, compresslevel=9, fileobj=None): 16 gzip.GzipFile.__init__(self, filename, mode, compresslevel, fileobj) 17 18 ### Context manager (stolen from Python 2.7)### 19 def __enter__(self): 20 """Context management protocol. Returns self.""" 21 return self 22 23 def __exit__(self, *args): 24 """Context management protocol. Calls close()""" 25 self.close() 26 27 28def open(filename, mode="rb", compresslevel=9): 29 if hasattr(gzip.GzipFile, "__enter__"): 30 return gzip.open(filename, mode, compresslevel) 31 else: 32 return GzipFile(filename, mode, compresslevel) 33 34 35def open_fileobj(fileobj, mode="rb", compresslevel=9): 36 if hasattr(gzip.GzipFile, "__enter__"): 37 return gzip.GzipFile( 38 filename="", mode=mode, fileobj=fileobj, compresslevel=compresslevel 39 ) 40 return GzipFile( 41 filename="", mode=mode, fileobj=fileobj, compresslevel=compresslevel 42 ) 43 44 45def compress(data, compresslevel=9): 46 """ 47 Returns the data compressed at gzip level compression. 48 """ 49 buf = io.BytesIO() 50 with open_fileobj(buf, "wb", compresslevel) as ogz: 51 if not isinstance(data, bytes): 52 data = data.encode(__salt_system_encoding__) 53 ogz.write(data) 54 compressed = buf.getvalue() 55 return compressed 56 57 58def uncompress(data): 59 buf = io.BytesIO(data) 60 with open_fileobj(buf, "rb") as igz: 61 unc = igz.read() 62 return unc 63 64 65def compress_file(fh_, compresslevel=9, chunk_size=1048576): 66 """ 67 Generator that reads chunk_size bytes at a time from a file/filehandle and 68 yields the compressed result of each read. 69 70 .. note:: 71 Each chunk is compressed separately. They cannot be stitched together 72 to form a compressed file. This function is designed to break up a file 73 into compressed chunks for transport and decompression/reassembly on a 74 remote host. 75 """ 76 try: 77 bytes_read = int(chunk_size) 78 if bytes_read != chunk_size: 79 raise ValueError 80 except ValueError: 81 raise ValueError("chunk_size must be an integer") 82 try: 83 while bytes_read == chunk_size: 84 buf = io.BytesIO() 85 with open_fileobj(buf, "wb", compresslevel) as ogz: 86 try: 87 bytes_read = ogz.write(fh_.read(chunk_size)) 88 except AttributeError: 89 # Open the file and re-attempt the read 90 fh_ = salt.utils.files.fopen(fh_, "rb") 91 bytes_read = ogz.write(fh_.read(chunk_size)) 92 yield buf.getvalue() 93 finally: 94 try: 95 fh_.close() 96 except AttributeError: 97 pass 98