summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rwxr-xr-xbitbake/bin/bitbake-selftest1
-rw-r--r--bitbake/lib/bb/compress/_pipecompress.py194
-rw-r--r--bitbake/lib/bb/compress/lz4.py17
-rw-r--r--bitbake/lib/bb/compress/zstd.py28
-rw-r--r--bitbake/lib/bb/tests/compression.py98
5 files changed, 338 insertions, 0 deletions
diff --git a/bitbake/bin/bitbake-selftest b/bitbake/bin/bitbake-selftest
index 6c0737416b..aec4706921 100755
--- a/bitbake/bin/bitbake-selftest
+++ b/bitbake/bin/bitbake-selftest
@@ -29,6 +29,7 @@ tests = ["bb.tests.codeparser",
29 "bb.tests.runqueue", 29 "bb.tests.runqueue",
30 "bb.tests.siggen", 30 "bb.tests.siggen",
31 "bb.tests.utils", 31 "bb.tests.utils",
32 "bb.tests.compression",
32 "hashserv.tests", 33 "hashserv.tests",
33 "layerindexlib.tests.layerindexobj", 34 "layerindexlib.tests.layerindexobj",
34 "layerindexlib.tests.restapi", 35 "layerindexlib.tests.restapi",
diff --git a/bitbake/lib/bb/compress/_pipecompress.py b/bitbake/lib/bb/compress/_pipecompress.py
new file mode 100644
index 0000000000..4b9f662143
--- /dev/null
+++ b/bitbake/lib/bb/compress/_pipecompress.py
@@ -0,0 +1,194 @@
1#
2# SPDX-License-Identifier: GPL-2.0-only
3#
4# Helper library to implement streaming compression and decompression using an
5# external process
6#
7# This library should be used directly by end users; a wrapper library for the
8# specific compression tool should be created
9
10import builtins
11import io
12import os
13import subprocess
14
15
16def open_wrap(
17 cls, filename, mode="rb", *, encoding=None, errors=None, newline=None, **kwargs
18):
19 """
20 Open a compressed file in binary or text mode.
21
22 Users should not call this directly. A specific compression library can use
23 this helper to provide it's own "open" command
24
25 The filename argument can be an actual filename (a str or bytes object), or
26 an existing file object to read from or write to.
27
28 The mode argument can be "r", "rb", "w", "wb", "x", "xb", "a" or "ab" for
29 binary mode, or "rt", "wt", "xt" or "at" for text mode. The default mode is
30 "rb".
31
32 For binary mode, this function is equivalent to the cls constructor:
33 cls(filename, mode). In this case, the encoding, errors and newline
34 arguments must not be provided.
35
36 For text mode, a cls object is created, and wrapped in an
37 io.TextIOWrapper instance with the specified encoding, error handling
38 behavior, and line ending(s).
39 """
40 if "t" in mode:
41 if "b" in mode:
42 raise ValueError("Invalid mode: %r" % (mode,))
43 else:
44 if encoding is not None:
45 raise ValueError("Argument 'encoding' not supported in binary mode")
46 if errors is not None:
47 raise ValueError("Argument 'errors' not supported in binary mode")
48 if newline is not None:
49 raise ValueError("Argument 'newline' not supported in binary mode")
50
51 file_mode = mode.replace("t", "")
52 if isinstance(filename, (str, bytes, os.PathLike)):
53 binary_file = cls(filename, file_mode, **kwargs)
54 elif hasattr(filename, "read") or hasattr(filename, "write"):
55 binary_file = cls(None, file_mode, fileobj=filename, **kwargs)
56 else:
57 raise TypeError("filename must be a str or bytes object, or a file")
58
59 if "t" in mode:
60 return io.TextIOWrapper(
61 binary_file, encoding, errors, newline, write_through=True
62 )
63 else:
64 return binary_file
65
66
67class CompressionError(OSError):
68 pass
69
70
71class PipeFile(io.RawIOBase):
72 """
73 Class that implements generically piping to/from a compression program
74
75 Derived classes should add the function get_compress() and get_decompress()
76 that return the required commands. Input will be piped into stdin and the
77 (de)compressed output should be written to stdout, e.g.:
78
79 class FooFile(PipeCompressionFile):
80 def get_decompress(self):
81 return ["fooc", "--decompress", "--stdout"]
82
83 def get_compress(self):
84 return ["fooc", "--compress", "--stdout"]
85
86 """
87
88 READ = 0
89 WRITE = 1
90
91 def __init__(self, filename=None, mode="rb", *, stderr=None, fileobj=None):
92 if "t" in mode or "U" in mode:
93 raise ValueError("Invalid mode: {!r}".format(mode))
94
95 if not "b" in mode:
96 mode += "b"
97
98 if mode.startswith("r"):
99 self.mode = self.READ
100 elif mode.startswith("w"):
101 self.mode = self.WRITE
102 else:
103 raise ValueError("Invalid mode %r" % mode)
104
105 if fileobj is not None:
106 self.fileobj = fileobj
107 else:
108 self.fileobj = builtins.open(filename, mode or "rb")
109
110 if self.mode == self.READ:
111 self.p = subprocess.Popen(
112 self.get_decompress(),
113 stdin=self.fileobj,
114 stdout=subprocess.PIPE,
115 stderr=stderr,
116 close_fds=True,
117 )
118 self.pipe = self.p.stdout
119 else:
120 self.p = subprocess.Popen(
121 self.get_compress(),
122 stdin=subprocess.PIPE,
123 stdout=self.fileobj,
124 stderr=stderr,
125 close_fds=True,
126 )
127 self.pipe = self.p.stdin
128
129 self.__closed = False
130
131 def _check_process(self):
132 if self.p is None:
133 return
134
135 returncode = self.p.wait()
136 if returncode:
137 raise CompressionError("Process died with %d" % returncode)
138 self.p = None
139
140 def close(self):
141 if self.closed:
142 return
143
144 self.pipe.close()
145 if self.p is not None:
146 self._check_process()
147 self.fileobj.close()
148
149 self.__closed = True
150
151 @property
152 def closed(self):
153 return self.__closed
154
155 def fileno(self):
156 return self.pipe.fileno()
157
158 def flush(self):
159 self.pipe.flush()
160
161 def isatty(self):
162 return self.pipe.isatty()
163
164 def readable(self):
165 return self.mode == self.READ
166
167 def writable(self):
168 return self.mode == self.WRITE
169
170 def readinto(self, b):
171 if self.mode != self.READ:
172 import errno
173
174 raise OSError(
175 errno.EBADF, "read() on write-only %s object" % self.__class__.__name__
176 )
177 size = self.pipe.readinto(b)
178 if size == 0:
179 self._check_process()
180 return size
181
182 def write(self, data):
183 if self.mode != self.WRITE:
184 import errno
185
186 raise OSError(
187 errno.EBADF, "write() on read-only %s object" % self.__class__.__name__
188 )
189 data = self.pipe.write(data)
190
191 if not data:
192 self._check_process()
193
194 return data
diff --git a/bitbake/lib/bb/compress/lz4.py b/bitbake/lib/bb/compress/lz4.py
new file mode 100644
index 0000000000..0f6bc51a5b
--- /dev/null
+++ b/bitbake/lib/bb/compress/lz4.py
@@ -0,0 +1,17 @@
1#
2# SPDX-License-Identifier: GPL-2.0-only
3#
4
5import bb.compress._pipecompress
6
7
8def open(*args, **kwargs):
9 return bb.compress._pipecompress.open_wrap(LZ4File, *args, **kwargs)
10
11
12class LZ4File(bb.compress._pipecompress.PipeFile):
13 def get_compress(self):
14 return ["lz4c", "-z", "-c"]
15
16 def get_decompress(self):
17 return ["lz4c", "-d", "-c"]
diff --git a/bitbake/lib/bb/compress/zstd.py b/bitbake/lib/bb/compress/zstd.py
new file mode 100644
index 0000000000..50c42133fb
--- /dev/null
+++ b/bitbake/lib/bb/compress/zstd.py
@@ -0,0 +1,28 @@
1#
2# SPDX-License-Identifier: GPL-2.0-only
3#
4
5import bb.compress._pipecompress
6import shutil
7
8
9def open(*args, **kwargs):
10 return bb.compress._pipecompress.open_wrap(ZstdFile, *args, **kwargs)
11
12
13class ZstdFile(bb.compress._pipecompress.PipeFile):
14 def __init__(self, *args, num_threads=1, compresslevel=3, **kwargs):
15 self.num_threads = num_threads
16 self.compresslevel = compresslevel
17 super().__init__(*args, **kwargs)
18
19 def _get_zstd(self):
20 if self.num_threads == 1 or not shutil.which("pzstd"):
21 return ["zstd"]
22 return ["pzstd", "-p", "%d" % self.num_threads]
23
24 def get_compress(self):
25 return self._get_zstd() + ["-c", "-%d" % self.compresslevel]
26
27 def get_decompress(self):
28 return self._get_zstd() + ["-d", "-c"]
diff --git a/bitbake/lib/bb/tests/compression.py b/bitbake/lib/bb/tests/compression.py
new file mode 100644
index 0000000000..d3ddf67f1c
--- /dev/null
+++ b/bitbake/lib/bb/tests/compression.py
@@ -0,0 +1,98 @@
1#
2# SPDX-License-Identifier: GPL-2.0-only
3#
4
5from pathlib import Path
6import bb.compress.lz4
7import bb.compress.zstd
8import contextlib
9import os
10import shutil
11import tempfile
12import unittest
13import subprocess
14
15
16class CompressionTests(object):
17 def setUp(self):
18 self._t = tempfile.TemporaryDirectory()
19 self.tmpdir = Path(self._t.name)
20 self.addCleanup(self._t.cleanup)
21
22 def _file_helper(self, mode_suffix, data):
23 tmp_file = self.tmpdir / "compressed"
24
25 with self.do_open(tmp_file, mode="w" + mode_suffix) as f:
26 f.write(data)
27
28 with self.do_open(tmp_file, mode="r" + mode_suffix) as f:
29 read_data = f.read()
30
31 self.assertEqual(read_data, data)
32
33 def test_text_file(self):
34 self._file_helper("t", "Hello")
35
36 def test_binary_file(self):
37 self._file_helper("b", "Hello".encode("utf-8"))
38
39 def _pipe_helper(self, mode_suffix, data):
40 rfd, wfd = os.pipe()
41 with open(rfd, "rb") as r, open(wfd, "wb") as w:
42 with self.do_open(r, mode="r" + mode_suffix) as decompress:
43 with self.do_open(w, mode="w" + mode_suffix) as compress:
44 compress.write(data)
45 read_data = decompress.read()
46
47 self.assertEqual(read_data, data)
48
49 def test_text_pipe(self):
50 self._pipe_helper("t", "Hello")
51
52 def test_binary_pipe(self):
53 self._pipe_helper("b", "Hello".encode("utf-8"))
54
55 def test_bad_decompress(self):
56 tmp_file = self.tmpdir / "compressed"
57 with tmp_file.open("wb") as f:
58 f.write(b"\x00")
59
60 with self.assertRaises(OSError):
61 with self.do_open(tmp_file, mode="rb", stderr=subprocess.DEVNULL) as f:
62 data = f.read()
63
64
65class LZ4Tests(CompressionTests, unittest.TestCase):
66 def setUp(self):
67 if shutil.which("lz4c") is None:
68 self.skipTest("'lz4c' not found")
69 super().setUp()
70
71 @contextlib.contextmanager
72 def do_open(self, *args, **kwargs):
73 with bb.compress.lz4.open(*args, **kwargs) as f:
74 yield f
75
76
77class ZStdTests(CompressionTests, unittest.TestCase):
78 def setUp(self):
79 if shutil.which("zstd") is None:
80 self.skipTest("'zstd' not found")
81 super().setUp()
82
83 @contextlib.contextmanager
84 def do_open(self, *args, **kwargs):
85 with bb.compress.zstd.open(*args, **kwargs) as f:
86 yield f
87
88
89class PZStdTests(CompressionTests, unittest.TestCase):
90 def setUp(self):
91 if shutil.which("pzstd") is None:
92 self.skipTest("'pzstd' not found")
93 super().setUp()
94
95 @contextlib.contextmanager
96 def do_open(self, *args, **kwargs):
97 with bb.compress.zstd.open(*args, num_threads=2, **kwargs) as f:
98 yield f