From 9cc7824eef0a7bb985c172adbce3c438b25b41a6 Mon Sep 17 00:00:00 2001 From: Joerg Behrmann Date: Mon, 24 Apr 2023 14:44:01 +0200 Subject: [PATCH] Implement __bool__ and __str__ for compression enum --- mkosi/__init__.py | 4 ++-- mkosi/config.py | 6 +++--- mkosi/util.py | 10 ++++++++-- tests/test_util.py | 31 +++++++++++++++++++++++++++++++ 4 files changed, 44 insertions(+), 7 deletions(-) diff --git a/mkosi/__init__.py b/mkosi/__init__.py index 32db2b312..bac62c030 100644 --- a/mkosi/__init__.py +++ b/mkosi/__init__.py @@ -789,7 +789,7 @@ def compress_output(config: MkosiConfig, src: Path, uid: int, gid: int) -> None: if not src.is_file(): return - if config.compress_output == Compression.none: + if not config.compress_output: # If we shan't compress, then at least make the output file sparse with complete_step(f"Digging holes into output file {src}…"): run(["fallocate", "--dig-holes", src], user=uid, group=gid) @@ -1156,7 +1156,7 @@ def print_summary(config: MkosiConfig) -> None: Output Signature: {none_to_na(config.output_signature if config.sign else None)} Output nspawn Settings: {none_to_na(config.output_nspawn_settings if config.nspawn_settings is not None else None)} Incremental: {yes_no(config.incremental)} - Compression: {config.compress_output.name} + Compression: {config.compress_output} Bootable: {config.bootable} Kernel Command Line: {" ".join(config.kernel_command_line)} UEFI SecureBoot: {yes_no(config.secure_boot)} diff --git a/mkosi/config.py b/mkosi/config.py index f8d9fe601..941e94b62 100644 --- a/mkosi/config.py +++ b/mkosi/config.py @@ -1687,10 +1687,10 @@ class MkosiConfig: @property def output_compressed(self) -> Path: - if self.compress_output == Compression.none: + if not self.compress_output: return self.output - return self.output.parent / f"{self.output.name}.{self.compress_output.name}" + return self.output.parent / f"{self.output.name}.{self.compress_output}" def output_paths(self) -> tuple[Path, ...]: return ( @@ -1941,7 +1941,7 @@ def load_args(args: argparse.Namespace) -> MkosiConfig: opname = "acquire shell" if args.verb == Verb.shell else "boot" if args.output_format in (OutputFormat.tar, OutputFormat.cpio): die(f"Sorry, can't {opname} with a {args.output_format} archive.") - if args.compress_output != Compression.none: + if args.compress_output: die(f"Sorry, can't {opname} with a compressed image.") if args.repo_dirs and not ( diff --git a/mkosi/util.py b/mkosi/util.py index dedf5728d..144c95795 100644 --- a/mkosi/util.py +++ b/mkosi/util.py @@ -83,8 +83,8 @@ class Distribution(enum.Enum): return self.name -class Compression(str, enum.Enum): - none = "none" +class Compression(enum.Enum): + none = None zst = "zst" xz = "xz" bz2 = "bz2" @@ -92,6 +92,12 @@ class Compression(str, enum.Enum): lz4 = "lz4" lzma = "lzma" + def __str__(self) -> str: + return str(self.value).lower() + + def __bool__(self) -> bool: + return bool(self.value) + def dictify(f: Callable[..., Iterator[tuple[T, V]]]) -> Callable[..., dict[T, V]]: def wrapper(*args: Any, **kwargs: Any) -> dict[T, V]: diff --git a/tests/test_util.py b/tests/test_util.py index 25df36848..c8da2dbbd 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -3,6 +3,7 @@ import os from mkosi.util import ( + Compression, Distribution, PackageType, set_umask, @@ -25,3 +26,33 @@ def test_set_umask() -> None: assert tmp1 == 0o767 assert tmp2 == 0o757 assert tmp3 == 0o777 + + +def test_compression_enum_creation() -> None: + assert Compression(None) == Compression.none + assert Compression("zst") == Compression.zst + assert Compression("xz") == Compression.xz + assert Compression("bz2") == Compression.bz2 + assert Compression("gz") == Compression.gz + assert Compression("lz4") == Compression.lz4 + assert Compression("lzma") == Compression.lzma + + +def test_compression_enum_bool() -> None: + assert bool(Compression.none) == False + assert bool(Compression.zst) == True + assert bool(Compression.xz) == True + assert bool(Compression.bz2) == True + assert bool(Compression.gz) == True + assert bool(Compression.lz4) == True + assert bool(Compression.lzma) == True + + +def test_compression_enum_str() -> None: + assert str(Compression.none) == "none" + assert str(Compression.zst) == "zst" + assert str(Compression.xz) == "xz" + assert str(Compression.bz2) == "bz2" + assert str(Compression.gz) == "gz" + assert str(Compression.lz4) == "lz4" + assert str(Compression.lzma) == "lzma" -- 2.47.2