xref: /qemu/tests/functional/qemu_test/uncompress.py (revision 4be0fce498d0a08f18b3a9accdb9ded79484d30a) 
1# SPDX-License-Identifier: GPL-2.0-or-later
2#
3# Utilities for python-based QEMU tests
4#
5# Copyright 2024 Red Hat, Inc.
6#
7# Authors:
8#  Thomas Huth <thuth@redhat.com>
9
10import gzip
11import lzma
12import os
13import shutil
14from urllib.parse import urlparse
15
16from .asset import Asset
17
18
19def gzip_uncompress(gz_path, output_path):
20    if os.path.exists(output_path):
21        return
22    with gzip.open(gz_path, 'rb') as gz_in:
23        try:
24            with open(output_path, 'wb') as raw_out:
25                shutil.copyfileobj(gz_in, raw_out)
26        except:
27            os.remove(output_path)
28            raise
29
30def lzma_uncompress(xz_path, output_path):
31    if os.path.exists(output_path):
32        return
33    with lzma.open(xz_path, 'rb') as lzma_in:
34        try:
35            with open(output_path, 'wb') as raw_out:
36                shutil.copyfileobj(lzma_in, raw_out)
37        except:
38            os.remove(output_path)
39            raise
40
41'''
42@params compressed: filename, Asset, or file-like object to uncompress
43@params uncompressed: filename to uncompress into
44@params format: optional compression format (gzip, lzma)
45
46Uncompresses @compressed into @uncompressed
47
48If @format is None, heuristics will be applied to guess the format
49from the filename or Asset URL. @format must be non-None if @uncompressed
50is a file-like object.
51
52Returns the fully qualified path to the uncompessed file
53'''
54def uncompress(compressed, uncompressed, format=None):
55    if format is None:
56        format = guess_uncompress_format(compressed)
57
58    if format == "xz":
59        lzma_uncompress(str(compressed), uncompressed)
60    elif format == "gz":
61        gzip_uncompress(str(compressed), uncompressed)
62    else:
63        raise Exception(f"Unknown compression format {format}")
64
65'''
66@params compressed: filename, Asset, or file-like object to guess
67
68Guess the format of @compressed, raising an exception if
69no format can be determined
70'''
71def guess_uncompress_format(compressed):
72    if type(compressed) == Asset:
73        compressed = urlparse(compressed.url).path
74    elif type(compressed) != str:
75        raise Exception(f"Unable to guess compression cformat for {compressed}")
76
77    (name, ext) = os.path.splitext(compressed)
78    if ext == ".xz":
79        return "xz"
80    elif ext == ".gz":
81        return "gz"
82    else:
83        raise Exception(f"Unknown compression format for {compressed}")
84