[PATCH v3 05/13] binman: Simplify comp_util class

[Stefan Herbrechtsmeier]

From: Stefan Herbrechtsmeier <stefan.herbrechtsmeier at weidmueller.com>

Simplify the comp_util class by remove duplicate code as preparation for
additional compress algorithms.

Signed-off-by: Stefan Herbrechtsmeier <stefan.herbrechtsmeier at weidmueller.com>

---

Changes in v3:
- Rename COMPRESSIONS to ALGORITHMS
- Rework existing comments
- Add _get_tool_name function comment
- Add _get_tool function comment

 tools/binman/cbfs_util_test.py |   2 +-
 tools/binman/comp_util.py      | 101 +++++++++++++++++++++++----------
 tools/binman/ftest.py          |   2 +-
 3 files changed, 72 insertions(+), 33 deletions(-)

diff --git a/tools/binman/cbfs_util_test.py b/tools/binman/cbfs_util_test.py
index f86b295149..44ebd04278 100755
--- a/tools/binman/cbfs_util_test.py
+++ b/tools/binman/cbfs_util_test.py
@@ -50,7 +50,7 @@ class TestCbfs(unittest.TestCase):
         cls.cbfstool = bintool.Bintool.create('cbfstool')
         cls.have_cbfstool = cls.cbfstool.is_present()
 
-        cls.have_lz4 = comp_util.HAVE_LZ4
+        cls.have_lz4 = comp_util.is_present('lz4')
 
     @classmethod
     def tearDownClass(cls):
diff --git a/tools/binman/comp_util.py b/tools/binman/comp_util.py
index 269bbf7975..00761d44cc 100644
--- a/tools/binman/comp_util.py
+++ b/tools/binman/comp_util.py
@@ -1,69 +1,108 @@
 # SPDX-License-Identifier: GPL-2.0+
 # Copyright 2022 Google LLC
+# Copyright (C) 2022 Weidmüller Interface GmbH & Co. KG
+# Stefan Herbrechtsmeier <stefan.herbrechtsmeier at weidmueller.com>
 #
-"""Utilities to compress and decompress data"""
+"""Utilities to compress and decompress data
+
+This supports the following compression algorithm:
+  none
+  lz4
+  lzma
+
+Note that for lzma this uses an old version of the algorithm, not that
+provided by xz.
+
+This requires the following tools:
+  lz4
+  lzma_alone
+
+It also requires an output directory to be previously set up, by calling
+PrepareOutputDir().
+"""
 
 import tempfile
 
 from binman import bintool
 from patman import tools
 
-LZ4 = bintool.Bintool.create('lz4')
-HAVE_LZ4 = LZ4.is_present()
+# Supported compression algorithms
+ALGORITHMS = ['lz4', 'lzma']
 
-LZMA_ALONE = bintool.Bintool.create('lzma_alone')
-HAVE_LZMA_ALONE = LZMA_ALONE.is_present()
+bintools = {}
 
+def _get_tool_name(algo):
+    """Get the tool name of a compression algorithm
 
-def compress(indata, algo):
-    """Compress some data using a given algorithm
+    Args:
+        algo (str): Algorithm to use
 
-    Note that for lzma this uses an old version of the algorithm, not that
-    provided by xz.
+    Returns:
+        str: Tool name
+    """
+    names = {'lzma': 'lzma_alone'}
+    return names.get(algo, algo)
+
+def _get_tool(algo):
+    """Get the bintool object of a compression algorithm
 
-    This requires 'lz4' and 'lzma_alone' tools. It also requires an output
-    directory to be previously set up, by calling PrepareOutputDir().
+    The function creates new bintool object on demand per compression algorithm
+    and save it in a global bintools dictionary.
+
+    Args:
+        algo (str): Algorithm to use
+
+    Returns:
+        A bintool object for the compression algorithm
+    """
+    global bintools
+    name = _get_tool_name(algo)
+    tool = bintools.get(name)
+    if not tool:
+        tool = bintool.Bintool.create(name)
+        bintools[name] = tool
+    return tool
+
+def compress(indata, algo):
+    """Compress some data using a given algorithm
 
     Args:
         indata (bytes): Input data to compress
-        algo (str): Algorithm to use ('none', 'lz4' or 'lzma')
+        algo (str): Algorithm to use
 
     Returns:
         bytes: Compressed data
     """
     if algo == 'none':
         return indata
-    if algo == 'lz4':
-        data = LZ4.compress(indata)
-    # cbfstool uses a very old version of lzma
-    elif algo == 'lzma':
-        data = LZMA_ALONE.compress(indata)
-    else:
+    if algo not in ALGORITHMS:
         raise ValueError("Unknown algorithm '%s'" % algo)
+
+    tool = _get_tool(algo)
+    data = tool.compress(indata)
+
     return data
 
 def decompress(indata, algo):
     """Decompress some data using a given algorithm
 
-    Note that for lzma this uses an old version of the algorithm, not that
-    provided by xz.
-
-    This requires 'lz4' and 'lzma_alone' tools. It also requires an output
-    directory to be previously set up, by calling PrepareOutputDir().
-
     Args:
         indata (bytes): Input data to decompress
-        algo (str): Algorithm to use ('none', 'lz4' or 'lzma')
+        algo (str): Algorithm to use
 
     Returns:
-        (bytes) Compressed data
+        bytes: Decompressed data
     """
     if algo == 'none':
         return indata
-    if algo == 'lz4':
-        data = LZ4.decompress(indata)
-    elif algo == 'lzma':
-        data = LZMA_ALONE.decompress(indata)
-    else:
+    if algo not in ALGORITHMS:
         raise ValueError("Unknown algorithm '%s'" % algo)
+
+    tool = _get_tool(algo)
+    data = tool.decompress(indata)
+
     return data
+
+def is_present(algo):
+     tool = _get_tool(algo)
+     return tool.is_present()
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index 057b4e28b7..96c15cff77 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -212,7 +212,7 @@ class TestFunctional(unittest.TestCase):
         TestFunctional._MakeInputFile('tee.elf',
             tools.read_file(cls.ElfTestFile('elf_sections')))
 
-        cls.have_lz4 = comp_util.HAVE_LZ4
+        cls.have_lz4 = comp_util.is_present('lz4')
 
     @classmethod
     def tearDownClass(cls):
-- 
2.30.2