1#
2# Copyright BitBake Contributors
3#
4# 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", "")
54    if isinstance(filename, (str, bytes, os.PathLike, int)):
55        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
197