1import io 

2import os 

3import tarfile 

4import time 

5from typing import Generator, Tuple 

6from logging import getLogger 

7from .iface import Input, Output 

8 

9log = getLogger(__name__) 

10 

11 

12class TarInput(Input): 

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

14 super().__init__(ifn) 

15 self.tf = tarfile.open(ifn) 

16 self.encoding = "utf-8" 

17 

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

19 # does not work? 

20 while True: 

21 n = self.tf.next() 

22 log.info("next %s", n) 

23 if n is None: 

24 break 

25 if not n.isfile(): 

26 continue 

27 yield n.name, n.mode, n.mtime 

28 

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

30 for n in self.tf.getmembers(): 

31 if not n.isfile(): 

32 continue 

33 yield n.name, n.mode, n.mtime 

34 

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

36 return self.tf.extractfile(fn).read().decode(self.encoding) 

37 

38 

39class TarOutput(Output): 

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

41 super().__init__(ofn) 

42 base, ext = os.path.splitext(ofn) 

43 ext = ext[1:] 

44 if ext == "tar": 

45 ext = "" 

46 self.tf = tarfile.open(ofn, "w:{}".format(ext)) 

47 self.encoding = "utf-8" 

48 self.owner = "root" 

49 self.group = "root" 

50 

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

52 bcont = content.encode(self.encoding) 

53 tinfo = tarfile.TarInfo(fn) 

54 tinfo.size = len(bcont) 

55 tinfo.mode = mode 

56 tinfo.type = tarfile.REGTYPE 

57 tinfo.uname = self.owner 

58 tinfo.gname = self.group 

59 if ts is None: 

60 tinfo.mtime = time.time() 

61 else: 

62 tinfo.mtime = ts 

63 self.tf.addfile(tinfo, io.BytesIO(bcont)) 

64 

65 def finish(self): 

66 self.tf.close()