Coverage for src / c41811 / config / processor / tarfile.py: 100%

1# cython: language_level = 3 # noqa: ERA001 

2 

3 

4""" 

5Tar压缩配置文件处理器 

6 

7.. versionadded:: 0.2.0 

8""" 

9 

10import itertools 

11import os 

12import tarfile 

13from collections.abc import Callable 

14from dataclasses import dataclass 

15from enum import ReprEnum 

16from typing import Any 

17from typing import Literal 

18from typing import cast 

19from typing import override 

20 

21from ..basic.core import ConfigFile 

22from ..main import BasicCompressedConfigSL 

23from ..safe_writer import safe_open 

24 

25 

26@dataclass(frozen=True) 

27class TarCompressionType: 

28 """压缩类型数据结构""" 

29 

30 full_name: str 

31 short_name: str | None 

32 

33 

34class TarCompressionTypes(TarCompressionType, ReprEnum): 

35 """压缩类型""" 

36 

37 ONLY_STORAGE = ("only-storage", None) 

38 

39 GZIP = ("gzip", "gz") 

40 BZIP2 = ("bzip2", "bz2") 

41 LZMA = ("lzma", "xz") 

42 

43 

44type ExtractionFilter = ( 

45 Literal["fully_trusted", "tar", "data"] | Callable[[tarfile.TarInfo, str], tarfile.TarInfo | None] 

46) 

47 

48 

49class TarFileSL(BasicCompressedConfigSL): 

50 """tar格式处理器""" 

51 

52 def __init__( 

53 self, 

54 *, 

55 reg_alias: str | None = None, 

56 create_dir: bool = True, 

57 compression: TarCompressionTypes | str | None = TarCompressionTypes.ONLY_STORAGE, 

58 compress_level: Literal[0, 1, 2, 3, 4, 5, 6, 7, 8, 9] | int | None = None, 

59 extraction_filter: ExtractionFilter | None = "data", 

60 ): 

61 """ 

62 :param reg_alias: sl处理器注册别名 

63 :type reg_alias: str | None 

64 :param create_dir: 是否创建目录 

65 :type create_dir: bool 

66 :param compression: 压缩类型 

67 :type compression: TarCompressionTypes | str | None 

68 :param compress_level: 压缩等级 

69 :type compress_level: Literal[0, 1, 2, 3, 4, 5, 6, 7, 8, 9] | int | None 

70 :param extraction_filter: 解压过滤器 

71 :type extraction_filter: ExtractionFilter | None 

72 """ # noqa: D205 

73 super().__init__(reg_alias=reg_alias, create_dir=create_dir) 

74 

75 if compression is None: 

76 compression = TarCompressionTypes.ONLY_STORAGE 

77 elif isinstance(compression, str): 

78 for compression_type in TarCompressionTypes: 

79 if compression in (compression_type.full_name, compression_type.short_name): 

80 compression = compression_type 

81 break 

82 

83 self._compression: TarCompressionType = cast(TarCompressionTypes, compression) 

84 self._compress_level: int | None = compress_level 

85 self._extraction_filter: ExtractionFilter | None = extraction_filter 

86 self._short_name = "" if self._compression.short_name is None else self._compression.short_name 

87 

88 @property 

89 @override 

90 def processor_reg_name(self) -> str: 

91 return f"tarfile:{self._short_name}" 

92 

93 @property 

94 @override 

95 def namespace_suffix(self) -> str: 

96 safe_name = self.processor_reg_name.replace(":", "-") 

97 return os.path.join(super().namespace_suffix, f"${safe_name}~") 

98 

99 @property 

100 @override 

101 def supported_file_patterns(self) -> tuple[str, ...]: 

102 if self._compression.short_name is None: 

103 return (".tar",) 

104 return f".tar.{self._compression.short_name}", f".tar.{self._compression.full_name}" 

105 

106 supported_file_classes = [ConfigFile] # noqa: RUF012 

107 

108 @override 

109 def compress_file(self, file_path: str, extract_dir: str) -> None: 

110 kwargs: dict[str, Any] = {} 

111 if self._compress_level is not None: 

112 # noinspection SpellCheckingInspection 

113 kwargs["compresslevel"] = self._compress_level 

114 with ( 

115 safe_open(file_path, "wb") as file, 

116 tarfile.open( 

117 mode=cast(Literal["w:", "w:gz", "w:bz2", "w:xz"], f"w:{self._short_name}"), 

118 fileobj=file, 

119 **kwargs, 

120 ) as tar, 

121 ): 

122 for root, dirs, files in os.walk(extract_dir): 

123 for item in itertools.chain(dirs, files): 

124 path = os.path.normpath(os.path.join(root, item)) 

125 tar.add(path, arcname=os.path.relpath(path, extract_dir), recursive=False) 

126 

127 @override 

128 def extract_file(self, file_path: str, extract_dir: str) -> None: 

129 with ( 

130 safe_open(file_path, "rb") as file, 

131 tarfile.open( 

132 mode=cast(Literal["r:", "r:gz", "r:bz2", "r:xz"], f"r:{self._short_name}"), 

133 fileobj=file, 

134 ) as tar, 

135 ): 

136 # py3.12不传入filter会发出警告 https://peps.python.org/pep-0706/#defaults-and-their-configuration 

137 tar.extractall(extract_dir, filter=self._extraction_filter) # noqa: S202 

138 

139 

140__all__ = ( 

141 "TarCompressionType", 

142 "TarCompressionTypes", 

143 "TarFileSL", 

144)