blob: 4a403d62cf8a406ed8922ffa50d2c3ded2d15678 [file] [log] [blame]
Andrew Geissler5f350902021-07-23 13:09:54 -04001#
Patrick Williams92b42cb2022-09-03 06:53:57 -05002# Copyright BitBake Contributors
3#
Andrew Geissler5f350902021-07-23 13:09:54 -04004# SPDX-License-Identifier: GPL-2.0-only
5#
6# Helper library to implement streaming compression and decompression using an
7# external process
8#
9# This library should be used directly by end users; a wrapper library for the
10# specific compression tool should be created
11
12import builtins
13import io
14import os
15import subprocess
16
17
18def open_wrap(
19 cls, filename, mode="rb", *, encoding=None, errors=None, newline=None, **kwargs
20):
21 """
22 Open a compressed file in binary or text mode.
23
24 Users should not call this directly. A specific compression library can use
25 this helper to provide it's own "open" command
26
27 The filename argument can be an actual filename (a str or bytes object), or
28 an existing file object to read from or write to.
29
30 The mode argument can be "r", "rb", "w", "wb", "x", "xb", "a" or "ab" for
31 binary mode, or "rt", "wt", "xt" or "at" for text mode. The default mode is
32 "rb".
33
34 For binary mode, this function is equivalent to the cls constructor:
35 cls(filename, mode). In this case, the encoding, errors and newline
36 arguments must not be provided.
37
38 For text mode, a cls object is created, and wrapped in an
39 io.TextIOWrapper instance with the specified encoding, error handling
40 behavior, and line ending(s).
41 """
42 if "t" in mode:
43 if "b" in mode:
44 raise ValueError("Invalid mode: %r" % (mode,))
45 else:
46 if encoding is not None:
47 raise ValueError("Argument 'encoding' not supported in binary mode")
48 if errors is not None:
49 raise ValueError("Argument 'errors' not supported in binary mode")
50 if newline is not None:
51 raise ValueError("Argument 'newline' not supported in binary mode")
52
53 file_mode = mode.replace("t", "")
Andrew Geisslereff27472021-10-29 15:35:00 -050054 if isinstance(filename, (str, bytes, os.PathLike, int)):
Andrew Geissler5f350902021-07-23 13:09:54 -040055 binary_file = cls(filename, file_mode, **kwargs)
56 elif hasattr(filename, "read") or hasattr(filename, "write"):
57 binary_file = cls(None, file_mode, fileobj=filename, **kwargs)
58 else:
59 raise TypeError("filename must be a str or bytes object, or a file")
60
61 if "t" in mode:
62 return io.TextIOWrapper(
63 binary_file, encoding, errors, newline, write_through=True
64 )
65 else:
66 return binary_file
67
68
69class CompressionError(OSError):
70 pass
71
72
73class PipeFile(io.RawIOBase):
74 """
75 Class that implements generically piping to/from a compression program
76
77 Derived classes should add the function get_compress() and get_decompress()
78 that return the required commands. Input will be piped into stdin and the
79 (de)compressed output should be written to stdout, e.g.:
80
81 class FooFile(PipeCompressionFile):
82 def get_decompress(self):
83 return ["fooc", "--decompress", "--stdout"]
84
85 def get_compress(self):
86 return ["fooc", "--compress", "--stdout"]
87
88 """
89
90 READ = 0
91 WRITE = 1
92
93 def __init__(self, filename=None, mode="rb", *, stderr=None, fileobj=None):
94 if "t" in mode or "U" in mode:
95 raise ValueError("Invalid mode: {!r}".format(mode))
96
97 if not "b" in mode:
98 mode += "b"
99
100 if mode.startswith("r"):
101 self.mode = self.READ
102 elif mode.startswith("w"):
103 self.mode = self.WRITE
104 else:
105 raise ValueError("Invalid mode %r" % mode)
106
107 if fileobj is not None:
108 self.fileobj = fileobj
109 else:
110 self.fileobj = builtins.open(filename, mode or "rb")
111
112 if self.mode == self.READ:
113 self.p = subprocess.Popen(
114 self.get_decompress(),
115 stdin=self.fileobj,
116 stdout=subprocess.PIPE,
117 stderr=stderr,
118 close_fds=True,
119 )
120 self.pipe = self.p.stdout
121 else:
122 self.p = subprocess.Popen(
123 self.get_compress(),
124 stdin=subprocess.PIPE,
125 stdout=self.fileobj,
126 stderr=stderr,
127 close_fds=True,
128 )
129 self.pipe = self.p.stdin
130
131 self.__closed = False
132
133 def _check_process(self):
134 if self.p is None:
135 return
136
137 returncode = self.p.wait()
138 if returncode:
139 raise CompressionError("Process died with %d" % returncode)
140 self.p = None
141
142 def close(self):
143 if self.closed:
144 return
145
146 self.pipe.close()
147 if self.p is not None:
148 self._check_process()
149 self.fileobj.close()
150
151 self.__closed = True
152
153 @property
154 def closed(self):
155 return self.__closed
156
157 def fileno(self):
158 return self.pipe.fileno()
159
160 def flush(self):
161 self.pipe.flush()
162
163 def isatty(self):
164 return self.pipe.isatty()
165
166 def readable(self):
167 return self.mode == self.READ
168
169 def writable(self):
170 return self.mode == self.WRITE
171
172 def readinto(self, b):
173 if self.mode != self.READ:
174 import errno
175
176 raise OSError(
177 errno.EBADF, "read() on write-only %s object" % self.__class__.__name__
178 )
179 size = self.pipe.readinto(b)
180 if size == 0:
181 self._check_process()
182 return size
183
184 def write(self, data):
185 if self.mode != self.WRITE:
186 import errno
187
188 raise OSError(
189 errno.EBADF, "write() on read-only %s object" % self.__class__.__name__
190 )
191 data = self.pipe.write(data)
192
193 if not data:
194 self._check_process()
195
196 return data