Coverage for src/epublib/source.py: 88%

1from datetime import datetime 

2from pathlib import Path 

3from typing import IO, Literal, Protocol, override 

4from zipfile import ZipInfo 

5 

6from epublib.exceptions import EPUBError 

7 

8 

9def zip_info_now(): 

10 now = datetime.now() 

11 

12 if now.year > 2107: 

13 now = now.replace(year=2107, month=12, day=31) 

14 

15 return ( 

16 now.year, 

17 now.month, 

18 now.day, 

19 now.hour, 

20 now.minute, 

21 now.second, 

22 ) 

23 

24 

25class SourceProtocol(Protocol): 

26 """Protocol for a source of EPUB data.""" 

27 

28 def getinfo(self, name: str) -> ZipInfo: ... 

29 def open( 

30 self, 

31 name: str | ZipInfo, 

32 mode: Literal["r", "w"] = "r", 

33 pwd: bytes | None = None, 

34 *, 

35 force_zip64: bool = False, 

36 ) -> IO[bytes]: ... 

37 def infolist(self) -> list[ZipInfo]: ... 

38 

39 

40class SinkProtocol(Protocol): 

41 """Protocol for a sink of EPUB data.""" 

42 

43 def writestr( 

44 self, 

45 zinfo_or_arcname: str | ZipInfo, 

46 data: bytes | str, 

47 compress_type: int | None = None, 

48 compresslevel: int | None = None, 

49 ) -> None: ... 

50 

51 

52class DirectorySource(SourceProtocol): 

53 """An EPUB source that reads the book from a directory on the filesystem (an 'unzipped' EPUB).""" 

54 

55 def __init__(self, path: str | Path): 

56 self.path: Path = Path(path) 

57 if not self.path.is_dir(): 

58 raise EPUBError(f"'{path}' is not a directory") 

59 

60 @override 

61 def infolist(self) -> list[ZipInfo]: 

62 return [ 

63 self._to_zipinfo(str(file.relative_to(self.path))) 

64 for file in self.path.rglob("*") 

65 if not file.is_dir() 

66 ] 

67 

68 @override 

69 def getinfo(self, name: str) -> ZipInfo: 

70 target = self.path / name 

71 if target.is_file(): 

72 return self._to_zipinfo(name) 

73 raise KeyError("There is no item named '{name}' in the folder") 

74 

75 def _to_zipinfo(self, name: str): 

76 return ZipInfo.from_file( 

77 self.path / name, 

78 arcname=name, 

79 strict_timestamps=False, 

80 ) 

81 

82 @override 

83 def open( 

84 self, 

85 name: str | ZipInfo, 

86 mode: Literal["r", "w"] = "r", 

87 pwd: bytes | None = None, 

88 *, 

89 force_zip64: bool = False, 

90 ) -> IO[bytes]: 

91 if isinstance(name, ZipInfo): 

92 filename = self.path / name.filename 

93 else: 

94 filename = self.path / name 

95 

96 return open(filename, mode + "b") 

97 

98 

99class DirectorySink(SinkProtocol): 

100 """An EPUB sink that writes the book to a directory on the filesystem (as an 'unzipped' EPUB).""" 

101 

102 def __init__(self, path: str | Path) -> None: 

103 self.path: Path = Path(path) 

104 if not self.path.is_dir(): 

105 raise EPUBError(f"'{path}' is not a directory") 

106 

107 @override 

108 def writestr( 

109 self, 

110 zinfo_or_arcname: str | ZipInfo, 

111 data: bytes | str, 

112 compress_type: int | None = None, 

113 compresslevel: int | None = None, 

114 ) -> None: 

115 if isinstance(zinfo_or_arcname, ZipInfo): 

116 filename = zinfo_or_arcname.filename 

117 else: 

118 filename = zinfo_or_arcname 

119 

120 full_path = self.path / filename 

121 

122 if str(full_path.parent) != ".": 

123 full_path.parent.mkdir(exist_ok=True, parents=True) 

124 

125 mode = "w" if isinstance(data, str) else "wb" 

126 with open(full_path, mode) as f: 

127 __ = f.write(data)