Coverage for src/ensembl/utils/archive.py: 95%

1# See the NOTICE file distributed with this work for additional information 

2# regarding copyright ownership. 

3# 

4# Licensed under the Apache License, Version 2.0 (the "License"); 

5# you may not use this file except in compliance with the License. 

6# You may obtain a copy of the License at 

7# 

8# http://www.apache.org/licenses/LICENSE-2.0 

9# 

10# Unless required by applicable law or agreed to in writing, software 

11# distributed under the License is distributed on an "AS IS" BASIS, 

12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

13# See the License for the specific language governing permissions and 

14# limitations under the License. 

15"""Utils for common IO operations over archive files, e.g. tar or gzip.""" 

16 

17from __future__ import annotations 

18 

19__all__ = [ 

20 "SUPPORTED_ARCHIVE_FORMATS", 

21 "open_gz_file", 

22 "extract_file", 

23] 

24 

25from contextlib import contextmanager 

26import gzip 

27from pathlib import Path 

28import shutil 

29import sys 

30from typing import Any, Generator, IO 

31 

32import ensembl.utils 

33from ensembl.utils import StrPath 

34from ensembl.utils.argparse import ArgumentParser 

35 

36 

37def _unpack_gz_files( 

38 src_file: StrPath, 

39 dst_dir: StrPath, 

40 **_kwargs: Any, 

41) -> None: 

42 """Unpacks `src_file` to `dst_dir`. 

43 

44 Args: 

45 src_file: File path to unpack (with ".gz" extension). 

46 dst_dir: Directory path to unpack the file into. 

47 

48 """ 

49 # Remove '.gz' extension to create the destination file name 

50 dst_file = Path(dst_dir) / Path(src_file).stem 

51 with gzip.open(src_file, "rb") as f_in: 

52 with dst_file.open("wb") as f_out: 

53 shutil.copyfileobj(f_in, f_out) 

54 

55 

56shutil.register_unpack_format("gzip", [".gz"], _unpack_gz_files, description="GZIP file") 

57 

58# Each registered format is a tuple (name, extensions, description) 

59SUPPORTED_ARCHIVE_FORMATS = [ext for elem in shutil.get_unpack_formats() for ext in elem[1]] 

60 

61 

62@contextmanager 

63def open_gz_file( 

64 file_path: StrPath, mode: str = "rt", encoding: str = "utf-8" 

65) -> Generator[gzip.GzipFile | IO, None, None]: 

66 """Yields an open file object, even if the file is compressed with gzip. 

67 

68 The file is expected to contain a text, and this can be used with the usual "with". 

69 

70 Args: 

71 file_path: A (single) file path to open. 

72 mode: The mode in which the file is opened. 

73 encoding: The name of the encoding used to decode or encode the file. 

74 

75 """ 

76 src_file = Path(file_path) 

77 if src_file.suffix == ".gz": 

78 with gzip.open(src_file, mode, encoding=encoding) as fh: 

79 yield fh 

80 else: 

81 with src_file.open(mode, encoding=encoding) as fh: 

82 yield fh 

83 

84 

85def extract_file(src_file: StrPath, dst_dir: StrPath) -> None: 

86 """Extracts the `src_file` into `dst_dir`. 

87 

88 If the file is not an archive, it will be copied to `dst_dir`. `dst_dir` will be created if it 

89 does not exist. 

90 

91 Args: 

92 src_file: Path to the file to unpack. 

93 dst_dir: Path to the folder where to extract the file. 

94 

95 """ 

96 src_file = Path(src_file) 

97 extensions = {"".join(src_file.suffixes[i:]) for i in range(0, len(src_file.suffixes))} 

98 

99 if extensions.intersection(SUPPORTED_ARCHIVE_FORMATS): 

100 if sys.version_info >= (3, 12): 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true

101 shutil.unpack_archive(src_file, dst_dir, filter="data") 

102 else: 

103 shutil.unpack_archive(src_file, dst_dir) 

104 else: 

105 # Replicate the functionality of shutil.unpack_archive() by creating `dst_dir` 

106 Path(dst_dir).mkdir(parents=True, exist_ok=True) 

107 shutil.copy(src_file, dst_dir) 

108 

109 

110def extract_file_cli() -> None: 

111 """Entry-point for the `extract_file` method""" 

112 parser = ArgumentParser(description="Extracts file to the given location.") 

113 parser.add_argument_src_path("--src_file", required=True, help="Path to the file to unpack") 

114 parser.add_argument_dst_path( 

115 "--dst_dir", default=Path.cwd(), help="Path to the folder where to extract the file" 

116 ) 

117 parser.add_argument("--version", action="version", version=ensembl.utils.__version__) 

118 args = parser.parse_args() 

119 extract_file(args.src_file, args.dst_dir)