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

1"""Zip archive utilities. 

2 

3Thin wrappers around ``zipfile`` that create and unpack zip archives in one call. 

4 

5Archives are created from a list of sources (``zip_to_path``, ``zip_to_bytes``). A source is a file path, 

6a directory path, which is scanned recursively, or a dict of archive names and contents. Entries are compressed 

7with ``ZIP_DEFLATED``. Archive names are the normalized source paths, optionally with ``base_dir`` stripped, 

8or only the base names with ``flat=True``. 

9 

10Archives are unpacked from a file or from bytes, into a directory (``unzip_path``, ``unzip_bytes``) 

11or into a dict of names and contents (``unzip_path_to_dict``, ``unzip_bytes_to_dict``). Directory entries 

12are skipped. Entries with unsafe names (absolute, starting with a dot or containing ``..``) are skipped 

13with a warning. With ``flat=True``, entries are unpacked by their base names, so entries with the same 

14base name overwrite each other. 

15 

16Example:: 

17 

18 import gws.lib.zipx as zipx 

19 

20 zipx.zip_to_path('/tmp/out.zip', ['/data/report', {'readme.txt': 'hello'}], base_dir='/data/') 

21 content = zipx.zip_to_bytes(['/data/a.txt', '/data/b.txt'], flat=True) 

22 

23 zipx.unzip_path('/tmp/out.zip', '/tmp/unpacked') 

24 files = zipx.unzip_bytes_to_dict(content) # {'a.txt': b'...', 'b.txt': b'...'} 

25""" 

26 

27import io 

28import os 

29import shutil 

30import zipfile 

31 

32import gws 

33 

34 

35class Error(gws.Error): 

36 """Zip archive error.""" 

37 

38 pass 

39 

40 

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

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

43 

44 If there are no files to add, no archive is created. 

45 

46 Args: 

47 path: Path to the archive. 

48 sources: File paths, directory paths (scanned recursively) or dicts that map 

49 archive names to contents (``str`` or ``bytes``). 

50 base_dir: Prefix to remove from the beginning of the file paths in the archive. 

51 flat: If ``True``, only the base names of the files are kept in the archive. 

52 

53 Returns: 

54 The number of files in the archive. 

55 

56 Raises: 

57 Error: If a source is neither a dict, a file nor a directory. 

58 """ 

59 

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

61 

62 

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

64 """Create a zip archive in memory. 

65 

66 Args: 

67 sources: File paths, directory paths (scanned recursively) or dicts that map 

68 archive names to contents (``str`` or ``bytes``). 

69 base_dir: Prefix to remove from the beginning of the file paths in the archive. 

70 flat: If ``True``, only the base names of the files are kept in the archive. 

71 

72 Returns: 

73 The archive content, or empty bytes if there are no files to add. 

74 

75 Raises: 

76 Error: If a source is neither a dict, a file nor a directory. 

77 """ 

78 

79 with io.BytesIO() as fp: 

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

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

82 

83 

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

85 """Unpack a zip archive file into a directory. 

86 

87 Args: 

88 path: Path to the archive. 

89 target_dir: Path to the target directory. 

90 flat: If ``True``, files are unpacked by their base names directly into ``target_dir``, 

91 which must exist; otherwise the directories of the archive are created as needed. 

92 

93 Returns: 

94 The number of unpacked files. 

95 """ 

96 

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

98 

99 

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

101 """Unpack a zip archive from bytes into a directory. 

102 

103 Args: 

104 source: The archive content. 

105 target_dir: Path to the target directory. 

106 flat: If ``True``, files are unpacked by their base names directly into ``target_dir``, 

107 which must exist; otherwise the directories of the archive are created as needed. 

108 

109 Returns: 

110 The number of unpacked files. 

111 """ 

112 

113 with io.BytesIO(source) as fp: 

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

115 

116 

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

118 """Unpack a zip archive file into a dict. 

119 

120 Args: 

121 path: Path to the archive. 

122 flat: If ``True``, the keys are the base names of the files, otherwise their paths in the archive. 

123 

124 Returns: 

125 A dict of file names and contents. 

126 """ 

127 

128 dct = {} 

129 _unzip(path, None, dct, flat) 

130 return dct 

131 

132 

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

134 """Unpack a zip archive from bytes into a dict. 

135 

136 Args: 

137 source: The archive content. 

138 flat: If ``True``, the keys are the base names of the files, otherwise their paths in the archive. 

139 

140 Returns: 

141 A dict of file names and contents. 

142 """ 

143 

144 with io.BytesIO(source) as fp: 

145 dct = {} 

146 _unzip(fp, None, dct, flat) 

147 return dct 

148 

149 

150## 

151 

152 

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

154 """Write sources to a zip file or file object, return the number of files.""" 

155 

156 def norm_path(p): 

157 p = os.path.normpath(p) 

158 if flat: 

159 return os.path.basename(p) 

160 if base_dir: 

161 if p.startswith(base_dir): 

162 return p[len(base_dir):] 

163 return p 

164 

165 def scan_dir(d): 

166 for de in os.scandir(d): 

167 if de.is_file(): 

168 yield de.path 

169 elif de.is_dir(): 

170 yield from scan_dir(de.path) 

171 

172 args = [] 

173 

174 for src in sources: 

175 if isinstance(src, dict): 

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

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

178 elif os.path.isdir(src): 

179 for p in scan_dir(src): 

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

181 elif os.path.isfile(src): 

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

183 else: 

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

185 

186 if not args: 

187 return 0 

188 

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

190 for arcname, path, data in args: 

191 if path: 

192 zf.write(path, arcname) 

193 else: 

194 zf.writestr(arcname, data) 

195 

196 return len(args) 

197 

198 

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

200 """Unpack a zip file or file object into a directory or a dict, return the number of files.""" 

201 

202 cnt = 0 

203 

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

205 for zi in zf.infolist(): 

206 if zi.is_dir(): 

207 continue 

208 

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

210 base = os.path.basename(path) 

211 

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

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

214 continue 

215 

216 cnt += 1 

217 

218 if target_dir: 

219 if flat: 

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

221 else: 

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

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

224 

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

226 shutil.copyfileobj(src, fp) 

227 elif target_dict is not None: 

228 key = base if flat else path 

229 with zf.open(zi) as src: 

230 target_dict[key] = src.read() 

231 else: 

232 raise Error('invalid target for unzip') 

233 

234 return cnt