diff --git a/mtda/client.py b/mtda/client.py index 47516771..c85a715d 100644 --- a/mtda/client.py +++ b/mtda/client.py @@ -245,6 +245,8 @@ def storage_write_image(self, path, callback=None): compression = CONSTS.IMAGE.GZ.value elif path.endswith(".zst"): compression = CONSTS.IMAGE.ZST.value + elif path.endswith(".xz"): + compression = CONSTS.IMAGE.XZ.value else: compression = CONSTS.IMAGE.RAW.value self._impl.storage_compression(compression, self._session) diff --git a/mtda/constants.py b/mtda/constants.py index dcd3b663..ba40998c 100644 --- a/mtda/constants.py +++ b/mtda/constants.py @@ -29,6 +29,7 @@ class IMAGE(Enum): BZ2 = 1 GZ = 2 ZST = 3 + XZ = 4 class MDNS: diff --git a/mtda/storage/writer.py b/mtda/storage/writer.py index 6f752719..d1d92213 100644 --- a/mtda/storage/writer.py +++ b/mtda/storage/writer.py @@ -15,6 +15,7 @@ import mtda.constants as CONSTS import zlib import zstandard as zstd +import lzma class AsyncImageWriter(queue.Queue): @@ -55,6 +56,8 @@ def compression(self, compression): self._write = self.write_gz elif compression == CONSTS.IMAGE.ZST: self._write = self.write_zst + elif compression == CONSTS.IMAGE.XZ: + self._write = self.write_xz else: raise ValueError("unsupported image compression!") self._compression = compression @@ -213,6 +216,32 @@ def write_zst(self, data): self.mtda.debug(3, "storage.writer.write_zst(): %s" % str(result)) return result + def write_xz(self, data): + self.mtda.debug(3, "storage.writer.write_xz()") + + # Create a xz decompressor when called for the first time + if self._zdec is None: + self._zdec = lzma.LZMADecompressor() + + try: + cont = True + result = None + while cont is True: + uncompressed = self._zdec.decompress(data, self._blksz) + result = self.storage.write(uncompressed) + self._written += result if result is not None else 0 + cont = self._zdec.needs_input is False + data = b'' + except EOFError: + result = 0 + except OSError as e: + self.mtda.debug(1, "storage.writer.write_xz(): " + "%s" % str(e.args[0])) + self._failed = True + + self.mtda.debug(3, "storage.writer.write_xz(): %s" % str(result)) + return result + @property def writing(self): return self._writing