Coverage for documents/utils.py: 33%

1import hashlib 

2import logging 

3import shutil 

4from collections.abc import Callable 

5from collections.abc import Iterable 

6from collections.abc import Iterator 

7from os import utime 

8from pathlib import Path 

9from subprocess import CompletedProcess 

10from subprocess import run 

11from typing import TYPE_CHECKING 

12from typing import Generic 

13from typing import TypeVar 

14 

15from django.conf import settings 

16from PIL import Image 

17 

18if TYPE_CHECKING: 18 ↛ 19line 18 didn't jump to line 19 because the condition on line 18 was never true

19 from django.db.models import Model 

20 from django.db.models import QuerySet 

21 

22_T = TypeVar("_T") 

23_M = TypeVar("_M", bound="Model") 

24 

25# A function that wraps an iterable — typically used to inject a progress bar. 

26IterWrapper = Callable[[Iterable[_T]], Iterable[_T]] 

27 

28 

29def identity(iterable: Iterable[_T]) -> Iterable[_T]: 

30 """Return the iterable unchanged; the no-op default for IterWrapper.""" 

31 return iterable 

32 

33 

34class QuerySetStream(Generic[_M]): 

35 """Stream a QuerySet via .iterator(chunk_size=...) instead of 

36 materializing it (plus any prefetch caches) all at once, while still 

37 supporting len() via count() so a progress bar wrapped around this 

38 (e.g. via IterWrapper) shows a real total instead of falling back to 

39 indeterminate. 

40 

41 Plain QuerySet iteration (``for row in queryset:``) is not lazy: Django 

42 fetches every matching row in one query and caches the fully-hydrated 

43 result in the queryset's own ``_result_cache`` before yielding the 

44 first item -- wrapping that in a progress bar or any other iterable 

45 adapter doesn't change this, since none of them alter how the 

46 underlying queryset produces items. ``.iterator(chunk_size=...)`` is 

47 the specific Django API that bypasses ``_result_cache`` and streams 

48 from a server-side cursor instead, discarding each chunk once consumed 

49 (and, since Django 4.1, still honours ``prefetch_related``, running the 

50 prefetches one batch at a time rather than for the whole queryset). 

51 

52 Subclass to layer additional per-batch work on top (see 

53 ``documents.search._backend._DocumentViewerStream``) by overriding 

54 ``__iter__`` -- ``__len__`` and the constructor are inherited for free. 

55 """ 

56 

57 def __init__(self, queryset: "QuerySet[_M]", *, chunk_size: int) -> None: 

58 self._queryset = queryset 

59 self._chunk_size = chunk_size 

60 

61 def __len__(self) -> int: 

62 return self._queryset.count() 

63 

64 def __iter__(self) -> Iterator[_M]: 

65 return iter(self._queryset.iterator(chunk_size=self._chunk_size)) 

66 

67 

68def _coerce_to_path( 

69 source: Path | str, 

70 dest: Path | str, 

71) -> tuple[Path, Path]: 

72 return Path(source).resolve(), Path(dest).resolve() 

73 

74 

75def copy_basic_file_stats(source: Path | str, dest: Path | str) -> None: 

76 """ 

77 Copies only the m_time and a_time attributes from source to destination. 

78 Both are expected to exist. 

79 

80 The extended attribute copy does weird things with SELinux and files 

81 copied from temporary directories and copystat doesn't allow disabling 

82 these copies. 

83 

84 If there is a PermissionError, skip copying file stats. 

85 """ 

86 source, dest = _coerce_to_path(source, dest) 

87 src_stat = source.stat() 

88 

89 try: 

90 utime(dest, ns=(src_stat.st_atime_ns, src_stat.st_mtime_ns)) 

91 except PermissionError: 

92 pass 

93 

94 

95def copy_file_with_basic_stats( 

96 source: Path | str, 

97 dest: Path | str, 

98) -> None: 

99 """ 

100 A sort of simpler copy2 that doesn't copy extended file attributes, 

101 only the access time and modified times from source to dest. 

102 

103 The extended attribute copy does weird things with SELinux and files 

104 copied from temporary directories. 

105 

106 If there is a PermissionError (e.g., on ZFS with acltype=nfsv4) 

107 fall back to copyfile (data only). 

108 """ 

109 source, dest = _coerce_to_path(source, dest) 

110 

111 try: 

112 shutil.copy(source, dest) 

113 except PermissionError: 

114 shutil.copyfile(source, dest) 

115 

116 copy_basic_file_stats(source, dest) 

117 

118 

119def maybe_override_pixel_limit() -> None: 

120 """ 

121 Maybe overrides the PIL limit on pixel count, if configured to allow it 

122 """ 

123 limit: float | int | None = settings.MAX_IMAGE_PIXELS 

124 if limit is not None and limit >= 0: 

125 pixel_count = limit 

126 if pixel_count == 0: 

127 pixel_count = None 

128 Image.MAX_IMAGE_PIXELS = pixel_count 

129 

130 

131def run_subprocess( 

132 arguments: list[str], 

133 env: dict[str, str] | None = None, 

134 logger: logging.Logger | None = None, 

135 *, 

136 check_exit_code: bool = True, 

137 log_stdout: bool = True, 

138 log_stderr: bool = True, 

139) -> CompletedProcess: 

140 """ 

141 Runs a subprocess and logs its output, checking return code if requested 

142 """ 

143 

144 proc_name = arguments[0] 

145 

146 completed_proc = run(args=arguments, env=env, capture_output=True, check=False) 

147 

148 if logger: 

149 logger.info(f"{proc_name} exited {completed_proc.returncode}") 

150 

151 if log_stdout and logger and completed_proc.stdout: 

152 stdout_str = ( 

153 completed_proc.stdout.decode("utf8", errors="ignore") 

154 .strip() 

155 .split( 

156 "\n", 

157 ) 

158 ) 

159 logger.info(f"{proc_name} stdout:") 

160 for line in stdout_str: 

161 logger.info(line) 

162 

163 if log_stderr and logger and completed_proc.stderr: 

164 stderr_str = ( 

165 completed_proc.stderr.decode("utf8", errors="ignore") 

166 .strip() 

167 .split( 

168 "\n", 

169 ) 

170 ) 

171 logger.info(f"{proc_name} stderr:") 

172 for line in stderr_str: 

173 logger.warning(line) 

174 

175 # Last, if requested, after logging outputs 

176 if check_exit_code: 

177 completed_proc.check_returncode() 

178 

179 return completed_proc 

180 

181 

182def get_boolean(boolstr: str) -> bool: 

183 """ 

184 Return a boolean value from a string representation. 

185 """ 

186 return bool(boolstr.lower() in ("yes", "y", "1", "t", "true")) 

187 

188 

189def compute_checksum(path: Path, chunk_size: int = 65536) -> str: 

190 """ 

191 Compute the SHA-256 checksum of a file. 

192 

193 Reads the file in chunks to avoid loading the entire file into memory. 

194 

195 Args: 

196 path (Path): Path to the file to hash. 

197 chunk_size (int, optional): Number of bytes to read per chunk. 

198 Defaults to 65536. 

199 

200 Returns: 

201 str: Hexadecimal SHA-256 digest of the file contents. 

202 

203 Raises: 

204 FileNotFoundError: If the file does not exist. 

205 OSError: If the file cannot be read. 

206 """ 

207 h = hashlib.sha256() 

208 with path.open("rb") as f: 

209 while chunk := f.read(chunk_size): 

210 h.update(chunk) 

211 return h.hexdigest()