1import datetime 

2import time 

3import zipfile 

4from typing import Generator, Tuple 

5from logging import getLogger 

6from .iface import Input, Output 

7 

8log = getLogger(__name__) 

9 

10 

11class ZipInput(Input): 

12 def __init__(self, ifn: str, **kwargs): 

13 super().__init__(ifn) 

14 self.zf = zipfile.ZipFile(ifn) 

15 self.encoding = "utf-8" 

16 self.ZIP_UNIX_SYSTEM = 3 

17 

18 def walk(self) -> Generator[Tuple[str, int, float], None, None]: 

19 for zinfo in self.zf.infolist(): 

20 if zinfo.is_dir(): 

21 continue 

22 ts = datetime.datetime(*zinfo.date_time) 

23 mode = 0o644 

24 if zinfo.create_system == self.ZIP_UNIX_SYSTEM: 

25 mode = zinfo.external_attr >> 16 

26 yield zinfo.filename, mode, ts.timestamp() 

27 

28 def readfile(self, fn: str) -> str: 

29 return self.zf.read(fn).decode(self.encoding) 

30 

31 

32class ZipOutput(Output): 

33 def __init__(self, ofn: str, **kwargs): 

34 super().__init__(ofn) 

35 self.zf = zipfile.ZipFile(ofn, "w") 

36 self.encoding = "utf-8" 

37 self.compress_type = zipfile.ZIP_DEFLATED 

38 

39 def writefile(self, fn: str, content: str, mode: int, ts: float = None): 

40 zinfo = zipfile.ZipInfo(fn, time.localtime(ts)) 

41 zinfo.compress_type = self.compress_type 

42 zinfo.external_attr = mode << 16 

43 if ts is not None: 

44 tm = time.localtime(ts) 

45 zinfo.date_time = (tm.tm_year, tm.tm_mon, tm.tm_mday, 

46 tm.tm_hour, tm.tm_min, tm.tm_sec) 

47 self.zf.writestr(zinfo, content) 

48 

49 def finish(self): 

50 self.zf.close()