Coverage for src/bioimageio/spec/_internal/io_basics.py: 96%

1from __future__ import annotations 

2 

3import hashlib 

4import os 

5import zipfile 

6from contextlib import nullcontext 

7from functools import partial 

8from pathlib import Path 

9from typing import Any, ClassVar, Protocol, runtime_checkable 

10from zipfile import ZipFile 

11 

12import pydantic 

13import zipp # pyright: ignore[reportMissingTypeStubs] 

14from annotated_types import Predicate 

15from pydantic import RootModel, StringConstraints 

16from typing_extensions import Annotated 

17 

18from .root_url import RootHttpUrl 

19from .validated_string import ValidatedString 

20 

21FileName = str 

22FilePath = Annotated[pydantic.FilePath, pydantic.Field(title="FilePath")] 

23AbsoluteDirectory = Annotated[ 

24 pydantic.DirectoryPath, 

25 Predicate(Path.is_absolute), 

26 pydantic.Field(title="AbsoluteDirectory"), 

27] 

28AbsoluteFilePath = Annotated[ 

29 pydantic.FilePath, 

30 Predicate(Path.is_absolute), 

31 pydantic.Field(title="AbsoluteFilePath"), 

32] 

33 

34BIOIMAGEIO_YAML = "rdf.yaml" 

35ALTERNATIVE_BIOIMAGEIO_YAML_NAMES = ("bioimageio.yaml", "model.yaml") 

36ALL_BIOIMAGEIO_YAML_NAMES = (BIOIMAGEIO_YAML,) + ALTERNATIVE_BIOIMAGEIO_YAML_NAMES 

37 

38ZipPath = zipp.Path # not zipfile.Path due to https://bugs.python.org/issue40564 

39 

40 

41class Sha256(ValidatedString): 

42 """A SHA-256 hash value""" 

43 

44 root_model: ClassVar[type[RootModel[Any]]] = RootModel[ 

45 Annotated[ 

46 str, 

47 StringConstraints( 

48 strip_whitespace=True, to_lower=True, min_length=64, max_length=64 

49 ), 

50 ] 

51 ] 

52 

53 

54class BytesReaderP(Protocol): 

55 def read(self, size: int = -1, /) -> bytes: ... 

56 

57 @property 

58 def closed(self) -> bool: ... 

59 

60 def readable(self) -> bool: ... 

61 

62 def seek(self, offset: int, whence: int = os.SEEK_SET, /) -> int: ... 

63 

64 def seekable(self) -> bool: ... 

65 

66 def tell(self) -> int: ... 

67 

68 

69@runtime_checkable 

70class BytesReaderIntoP(BytesReaderP, Protocol): 

71 def readinto(self, b: bytearray | memoryview) -> int: ... 

72 

73 

74Suffix = str 

75 

76 

77class BytesReader(BytesReaderP): 

78 def __init__( 

79 self, 

80 /, 

81 reader: BytesReaderP | BytesReaderIntoP, 

82 *, 

83 sha256: Sha256 | None, 

84 suffix: Suffix, 

85 original_file_name: FileName, 

86 original_root: RootHttpUrl | AbsoluteDirectory | ZipFile, 

87 is_zipfile: bool | None, 

88 ) -> None: 

89 self._reader = reader 

90 self._sha256 = sha256 

91 self._suffix = suffix 

92 self._original_file_name = original_file_name 

93 self._original_root = original_root 

94 self._is_zipfile = is_zipfile 

95 super().__init__() 

96 

97 @property 

98 def is_zipfile(self) -> bool: 

99 if self._is_zipfile is None: 

100 pos = self.tell() 

101 self._is_zipfile = zipfile.is_zipfile(self) 

102 _ = self.seek(pos) 

103 

104 return self._is_zipfile 

105 

106 @property 

107 def sha256(self) -> Sha256: 

108 if self._sha256 is None: 

109 pos = self._reader.tell() 

110 _ = self._reader.seek(0) 

111 self._sha256 = get_sha256(self._reader) 

112 _ = self._reader.seek(pos) 

113 

114 return self._sha256 

115 

116 @property 

117 def suffix(self) -> Suffix: 

118 return self._suffix 

119 

120 @property 

121 def original_file_name(self) -> FileName: 

122 return self._original_file_name 

123 

124 @property 

125 def original_root(self) -> RootHttpUrl | AbsoluteDirectory | ZipFile: 

126 return self._original_root 

127 

128 def read(self, size: int = -1, /) -> bytes: 

129 return self._reader.read(size) 

130 

131 def read_text(self, encoding: str = "utf-8") -> str: 

132 return self._reader.read().decode(encoding) 

133 

134 def readable(self) -> bool: 

135 return True 

136 

137 def seek(self, offset: int, whence: int = os.SEEK_SET, /) -> int: 

138 return self._reader.seek(offset, whence) 

139 

140 def seekable(self) -> bool: 

141 return True 

142 

143 def tell(self) -> int: 

144 return self._reader.tell() 

145 

146 @property 

147 def closed(self) -> bool: 

148 return self._reader.closed 

149 

150 

151def get_sha256(source: BytesReaderP | BytesReaderIntoP | Path) -> Sha256: 

152 chunksize = 128 * 1024 

153 h = hashlib.sha256() 

154 

155 if isinstance(source, BytesReaderIntoP): 

156 b = bytearray(chunksize) 

157 mv = memoryview(b) 

158 for n in iter(lambda: source.readinto(mv), 0): 

159 h.update(mv[:n]) 

160 else: 

161 if isinstance(source, Path): 

162 read_ctxt = source.open(mode="rb") 

163 else: 

164 read_ctxt = nullcontext(source) 

165 

166 with read_ctxt as r: 

167 for chunk in iter(partial(r.read, chunksize), b""): 

168 h.update(chunk) 

169 

170 sha = h.hexdigest() 

171 return Sha256(sha)