Newer
Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
# -*- coding: utf-8 -*-
import logging
import os
import tarfile
import tempfile
from pathlib import Path
from typing import Tuple
import zstandard
logger = logging.getLogger(__name__)
def decompress_zst_archive(compressed_archive: Path) -> Tuple[int, Path]:
"""
Decompress a ZST-compressed tar archive in data dir. The tar archive is not extracted.
This returns the path to the archive and the file descriptor.
Beware of closing the file descriptor explicitly or the main
process will keep the memory held even if the file is deleted.
:param compressed_archive: Path to the target ZST-compressed archive
:return: File descriptor and path to the uncompressed tar archive
"""
dctx = zstandard.ZstdDecompressor()
archive_fd, archive_path = tempfile.mkstemp(suffix=".tar")
logger.debug(f"Uncompressing file to {archive_path}")
try:
with open(compressed_archive, "rb") as compressed, open(
archive_path, "wb"
) as decompressed:
dctx.copy_stream(compressed, decompressed)
logger.debug(f"Successfully uncompressed archive {compressed_archive}")
except zstandard.ZstdError as e:
raise Exception(f"Couldn't uncompressed archive: {e}")
return archive_fd, Path(archive_path)
def extract_tar_archive(archive_path: Path, destination: Path):
"""
Extract the tar archive's content to a specific destination
:param archive_path: Path to the archive
:param destination: Path where the archive's data will be extracted
"""
try:
with tarfile.open(archive_path) as tar_archive:
tar_archive.extractall(destination)
except tarfile.ReadError as e:
raise Exception(f"Couldn't handle the decompressed Tar archive: {e}")
def extract_tar_zst_archive(
compressed_archive: Path, destination: Path
) -> Tuple[int, Path]:
"""
Extract a ZST-compressed tar archive's content to a specific destination
:param compressed_archive: Path to the target ZST-compressed archive
:param destination: Path where the archive's data will be extracted
:return: File descriptor and path to the uncompressed tar archive
"""
archive_fd, archive_path = decompress_zst_archive(compressed_archive)
extract_tar_archive(archive_path, destination)
return archive_fd, archive_path
def close_delete_file(file_descriptor: int, file_path: Path):
"""
Close the file descriptor of the file and delete the file
:param file_descriptor: File descriptor of the archive
:param file_path: Path to the archive
"""
try:
os.close(file_descriptor)
file_path.unlink()
except OSError as e:
logger.warning(f"Unable to delete file {file_path}: {e}")