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

1# cython: language_level = 3 # noqa: ERA001 

2 

3 

4""" 

5Zip压缩配置文件处理器 

6 

7.. versionadded:: 0.2.0 

8""" 

9 

10import itertools 

11import os 

12import zipfile 

13from dataclasses import dataclass 

14from enum import ReprEnum 

15from typing import Literal 

16from typing import cast 

17from typing import override 

18 

19from ..basic.core import ConfigFile 

20from ..main import BasicCompressedConfigSL 

21from ..safe_writer import safe_open 

22 

23 

24@dataclass(frozen=True) 

25class ZipCompressionType: 

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

27 

28 full_name: str 

29 short_name: str | None 

30 zipfile_constant: int 

31 

32 

33class ZipCompressionTypes(ZipCompressionType, ReprEnum): 

34 """压缩类型""" 

35 

36 ONLY_STORAGE = ("only-storage", None, zipfile.ZIP_STORED) 

37 

38 ZIP = ("zip", "zip", zipfile.ZIP_DEFLATED) 

39 BZIP2 = ("bzip2", "bz2", zipfile.ZIP_BZIP2) 

40 LZMA = ("lzma", "xz", zipfile.ZIP_LZMA) 

41 

42 

43class ZipFileSL(BasicCompressedConfigSL): 

44 """zip格式处理器""" 

45 

46 def __init__( 

47 self, 

48 *, 

49 reg_alias: str | None = None, 

50 create_dir: bool = True, 

51 compression: ZipCompressionTypes | str | int | None = ZipCompressionTypes.ONLY_STORAGE, 

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

53 ): 

54 """ 

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

56 :type reg_alias: str | None 

57 :param create_dir: 是否创建目录 

58 :type create_dir: bool 

59 :param compression: 压缩类型 

60 :type compression: ZipCompressionTypes | str | int | None 

61 :param compress_level: 压缩等级 

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

63 """ # noqa: D205 

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

65 

66 if compression is None: 

67 compression = ZipCompressionTypes.ONLY_STORAGE 

68 elif isinstance(compression, str | int): 

69 for compression_type in ZipCompressionTypes: 

70 if compression in ( 

71 compression_type.full_name, 

72 compression_type.short_name, 

73 compression_type.zipfile_constant, 

74 ): 

75 compression = compression_type 

76 break 

77 

78 self._compression: ZipCompressionType = cast(ZipCompressionTypes, compression) 

79 self._compress_level: int | None = compress_level 

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

81 

82 @property 

83 @override 

84 def processor_reg_name(self) -> str: 

85 return f"zipfile:{self._short_name}-{self._compress_level}" 

86 

87 @property 

88 @override 

89 def namespace_suffix(self) -> str: 

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

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

92 

93 @property 

94 @override 

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

96 if self._compression.short_name is None: 

97 return f".{self._compress_level}.zip", ".zip" 

98 return ( 

99 f".{self._compress_level}.{self._compression.full_name}", 

100 f".{self._compress_level}.{self._compression.short_name}", 

101 f".{self._compression.short_name}", 

102 f".{self._compression.full_name}", 

103 ) 

104 

105 supported_file_classes = [ConfigFile] # noqa: RUF012 

106 

107 @override 

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

109 with ( 

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

111 zipfile.ZipFile( 

112 file, mode="w", compression=self._compression.zipfile_constant, compresslevel=self._compress_level 

113 ) as zip_file, 

114 ): 

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

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

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

118 zip_file.write(path, arcname=os.path.relpath(path, extract_dir)) 

119 

120 @override 

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

122 with safe_open(file_path, "rb") as file, zipfile.ZipFile(file) as zip_file: 

123 zip_file.extractall(extract_dir) 

124 

125 

126__all__ = ( 

127 "ZipCompressionType", 

128 "ZipCompressionTypes", 

129 "ZipFileSL", 

130)