Coverage for gws-app/gws/lib/zipx/__init__.py: 95%

1"""Zipfile wrappers.""" 

2 

3import io 

4import os 

5import shutil 

6import zipfile 

7 

8import gws 

9 

10 

11class Error(gws.Error): 

12 pass 

13 

14 

15def zip_to_path(path: str, sources: list[str | dict], base_dir: str = '', flat: bool = False) -> int: 

16 """Create a zip archive in a file. 

17 

18 Args: 

19 path: Path to the archive. 

20 sources: A list of paths or dicts to zip. If a dict is given,  

21 its keys are file names in the archive and its values are the file contents. 

22 base_dir: If given, this path is stripped from the beginning of the file paths in the archive. 

23 flat: If ``True`` only base names are being kept in archive. 

24 

25 Returns: 

26 The amount of files in the archive. 

27 """ 

28 

29 return _zip(path, sources, base_dir, flat) 

30 

31 

32def zip_to_bytes(sources: list[str | dict], base_dir: str = '', flat: bool = False) -> bytes: 

33 """Create a zip archive in memory. 

34 

35 Args: 

36 sources: A list of paths or dicts to zip. If a dict is given,  

37 its keys are file names in the archive and its values are the file contents. 

38 base_dir: If given, this path is stripped from the beginning of the file paths in the archive. 

39 flat: If ``True`` only base names are being kept in archive. 

40 

41 Returns: 

42 The zipped content as bytes. 

43 """ 

44 

45 with io.BytesIO() as fp: 

46 cnt = _zip(fp, sources, base_dir, flat) 

47 return fp.getvalue() if cnt else b'' 

48 

49 

50def unzip_path(path: str, target_dir: str, flat: bool = False) -> int: 

51 """Unpack a zip archive into a directory. 

52 

53 Args: 

54 path: Path to the zip archive. 

55 target_dir: Path to the target directory. 

56 flat: If ``True`` omit path and consider only base name of files in the zip archive, 

57 else complete paths are considered of files in the zip archive. Default is ``False``. 

58 

59 Returns: 

60 The number of unzipped files. 

61 """ 

62 

63 return _unzip(path, target_dir, None, flat) 

64 

65 

66def unzip_bytes(source: bytes, target_dir: str, flat: bool = False) -> int: 

67 """Unpack a zip archive in memory into a directory. 

68 

69 Args: 

70 source: Path to the zip archive. 

71 target_dir: Path to the target directory. 

72 flat: If ``True`` omit path and consider only base name of files in the zip archive, 

73 else complete paths are considered of files in the zip archive. Default is ``False``. 

74 

75 Returns: 

76 The number of unzipped files. 

77 """ 

78 

79 with io.BytesIO(source) as fp: 

80 return _unzip(fp, target_dir, None, flat) 

81 

82 

83def unzip_path_to_dict(path: str, flat: bool = False) -> dict[str, bytes]: 

84 """Unpack a zip archive into a dict. 

85 

86 Args: 

87 path: Path to the zip archive. 

88 flat: If ``True`` then the result contains the base names of the unzipped files, 

89 else it contains the whole path. Default is ``False``. 

90 

91 Returns: 

92 A dictionary whose keys are the file paths or base names and values are the file contents. 

93 """ 

94 

95 dct = {} 

96 _unzip(path, None, dct, flat) 

97 return dct 

98 

99 

100def unzip_bytes_to_dict(source: bytes, flat: bool = False) -> dict[str, bytes]: 

101 """Unpack a zip archive in memory into a dict. 

102 

103 Args: 

104 source: Path to zip archive. 

105 flat: If ``True`` then the result contains the base names of the unzipped files, 

106 else it contains the whole path. Default is ``False``. 

107 

108 Returns: 

109 A dictionary whose keys are the file paths or base names and values are the file contents. 

110 """ 

111 

112 with io.BytesIO(source) as fp: 

113 dct = {} 

114 _unzip(fp, None, dct, flat) 

115 return dct 

116 

117 

118## 

119 

120 

121def _zip(target, sources, base_dir, flat): 

122 def norm_path(p): 

123 p = os.path.normpath(p) 

124 if flat: 

125 return os.path.basename(p) 

126 if base_dir: 

127 if p.startswith(base_dir): 

128 return p[len(base_dir):] 

129 return p 

130 

131 def scan_dir(d): 

132 for de in os.scandir(d): 

133 if de.is_file(): 

134 yield de.path 

135 elif de.is_dir(): 

136 yield from scan_dir(de.path) 

137 

138 args = [] 

139 

140 for src in sources: 

141 if isinstance(src, dict): 

142 for name, data in src.items(): 

143 args.append((norm_path(name), None, data)) 

144 elif os.path.isdir(src): 

145 for p in scan_dir(src): 

146 args.append((norm_path(p), p, None)) 

147 elif os.path.isfile(src): 

148 args.append((norm_path(src), src, None)) 

149 else: 

150 raise Error(f'zip: invalid argument: {src!r}') 

151 

152 if not args: 

153 return 0 

154 

155 with zipfile.ZipFile(target, 'w', compression=zipfile.ZIP_DEFLATED) as zf: 

156 for arcname, path, data in args: 

157 if path: 

158 zf.write(path, arcname) 

159 else: 

160 zf.writestr(arcname, data) 

161 

162 return len(args) 

163 

164 

165def _unzip(source, target_dir, target_dict, flat): 

166 cnt = 0 

167 

168 with zipfile.ZipFile(source, 'r') as zf: 

169 for zi in zf.infolist(): 

170 if zi.is_dir(): 

171 continue 

172 

173 path = zi.filename.replace('\\', '/') 

174 base = os.path.basename(path) 

175 

176 if path.startswith(('/', '.')) or '..' in path or not base: 

177 gws.log.warning(f'unzip: invalid file name: {path!r}') 

178 continue 

179 

180 cnt += 1 

181 

182 if target_dir: 

183 if flat: 

184 dst = os.path.join(target_dir, base) 

185 else: 

186 dst = os.path.join(target_dir, *path.split('/')) 

187 os.makedirs(os.path.dirname(dst), exist_ok=True) 

188 

189 with zf.open(zi) as src, open(dst, 'wb') as fp: 

190 shutil.copyfileobj(src, fp) 

191 elif target_dict is not None: 

192 key = base if flat else path 

193 with zf.open(zi) as src: 

194 target_dict[key] = src.read() 

195 else: 

196 raise Error('invalid target for unzip') 

197 

198 return cnt