Coverage for python/lsst/resources/utils.py: 76%

1# This file is part of lsst-resources. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# Use of this source code is governed by a 3-clause BSD-style 

10# license that can be found in the LICENSE file. 

11 

12from __future__ import annotations 

13 

14__all__ = ("NoTransaction", "TransactionProtocol", "get_tempdir", "os2posix", "posix2os") 

15 

16import contextlib 

17import logging 

18import multiprocessing 

19import os 

20import posixpath 

21import shutil 

22import stat 

23import tempfile 

24from collections.abc import Callable, Generator 

25from functools import cache 

26from pathlib import Path, PurePath, PurePosixPath 

27from typing import Any, Protocol 

28 

29# Determine if the path separator for the OS looks like POSIX 

30IS_POSIX = os.sep == posixpath.sep 

31 

32# Root path for this operating system. This can use getcwd which 

33# can fail in some situations so in the default case assume that 

34# posix means posix and only determine explicitly in the non-posix case. 

35OS_ROOT_PATH = posixpath.sep if IS_POSIX else Path().resolve().root 

36 

37# Default upper bound on the number of workers for parallelized operations. 

38# Subclasses of ResourcePath override this for schemes that have no connection 

39# pool to contend with. Backends that do have one size that pool from the 

40# worker count, so they need no separate coordination here. 

41MAX_WORKERS = 10 

42 

43log = logging.getLogger(__name__) 

44 

45 

46def os2posix(ospath: str) -> str: 

47 """Convert a local path description to a POSIX path description. 

48 

49 Parameters 

50 ---------- 

51 ospath : `str` 

52 Path using the local path separator. 

53 

54 Returns 

55 ------- 

56 posix : `str` 

57 Path using POSIX path separator. 

58 """ 

59 if IS_POSIX: 59 ↛ 62line 59 didn't jump to line 62 because the condition on line 59 was always true

60 return ospath 

61 

62 posix = PurePath(ospath).as_posix() 

63 

64 # PurePath strips trailing "/" from paths such that you can no 

65 # longer tell if a path is meant to be referring to a directory 

66 # Try to fix this. 

67 if ospath.endswith(os.sep) and not posix.endswith(posixpath.sep): 

68 posix += posixpath.sep 

69 

70 return posix 

71 

72 

73def posix2os(posix: PurePath | str) -> str: 

74 """Convert a POSIX path description to a local path description. 

75 

76 Parameters 

77 ---------- 

78 posix : `str`, `~pathlib.PurePath` 

79 Path using the POSIX path separator. 

80 

81 Returns 

82 ------- 

83 ospath : `str` 

84 Path using OS path separator. 

85 """ 

86 if IS_POSIX: 86 ↛ 89line 86 didn't jump to line 89 because the condition on line 86 was always true

87 return str(posix) 

88 

89 posixPath = PurePosixPath(posix) 

90 paths = list(posixPath.parts) 

91 

92 # Have to convert the root directory after splitting 

93 if paths[0] == posixPath.root: 

94 paths[0] = OS_ROOT_PATH 

95 

96 # Trailing "/" is stripped so we need to add back an empty path 

97 # for consistency 

98 if str(posix).endswith(posixpath.sep): 

99 paths.append("") 

100 

101 return os.path.join(*paths) 

102 

103 

104@cache 

105def get_tempdir() -> str: 

106 """Get POSIX path to temporary directory. 

107 

108 Returns 

109 ------- 

110 tmpdir : `str` 

111 Path to the default temporary directory location. 

112 

113 Notes 

114 ----- 

115 Uses the value of environment variables ``LSST_RESOURCES_TMPDIR`` or 

116 ``TMPDIR``, if defined. Otherwise use the system temporary directory, 

117 with a last-resort fallback to the current working directory if 

118 nothing else is available. 

119 """ 

120 tmpdir = None 

121 # $TMPDIR is also checked with getttempdir() below. 

122 for dir in (os.getenv(v) for v in ("LSST_RESOURCES_TMPDIR", "TMPDIR")): 

123 if dir and os.path.isdir(dir): 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true

124 tmpdir = dir 

125 break 

126 

127 if tmpdir is None: 127 ↛ 130line 127 didn't jump to line 130 because the condition on line 127 was always true

128 tmpdir = tempfile.gettempdir() 

129 

130 return tmpdir 

131 

132 

133class NoTransaction: 

134 """A simple emulation of the 

135 `~lsst.daf.butler.core.datastore.DatastoreTransaction` class. 

136 

137 Notes 

138 ----- 

139 Does nothing. Used as a fallback in the absence of an explicit transaction 

140 class. 

141 """ 

142 

143 def __init__(self) -> None: 

144 return 

145 

146 @contextlib.contextmanager 

147 def undoWith(self, name: str, undoFunc: Callable, *args: Any, **kwargs: Any) -> Generator[None]: 

148 """No-op context manager to replace 

149 `~lsst.daf.butler.core.datastore.DatastoreTransaction`. 

150 

151 Parameters 

152 ---------- 

153 name : `str` 

154 The name of this undo request. 

155 undoFunc : `~collections.abc.Callable` 

156 Function to call if there is an exception. Not used. 

157 *args : `~typing.Any` 

158 Parameters to pass to ``undoFunc``. 

159 **kwargs : `~typing.Any` 

160 Keyword parameters to pass to ``undoFunc``. 

161 

162 Yields 

163 ------ 

164 `None` 

165 Context manager returns nothing since transactions are disabled 

166 by definition. 

167 """ 

168 yield None 

169 

170 

171class TransactionProtocol(Protocol): 

172 """Protocol for type checking transaction interface.""" 

173 

174 @contextlib.contextmanager 

175 def undoWith(self, name: str, undoFunc: Callable, *args: Any, **kwargs: Any) -> Generator[None]: ... 175 ↛ exitline 175 didn't return from function 'undoWith' because

176 

177 

178def makeTestTempDir(default_base: str | None = None) -> str: 

179 """Create a temporary directory for test usage. 

180 

181 The directory will be created within ``LSST_RESOURCES_TEST_TMP`` if that 

182 environment variable is set, falling back to ``LSST_RESOURCES_TMPDIR`` 

183 amd then ``default_base`` if none are set. 

184 

185 Parameters 

186 ---------- 

187 default_base : `str`, optional 

188 Default parent directory. Will use system default if no environment 

189 variables are set and base is set to `None`. 

190 

191 Returns 

192 ------- 

193 dir : `str` 

194 Name of the new temporary directory. 

195 """ 

196 base = default_base 

197 for envvar in ("LSST_RESOURCES_TEST_TMP", "LSST_RESOURCES_TMPDIR"): 

198 if envvar in os.environ and os.environ[envvar]: 198 ↛ 199line 198 didn't jump to line 199 because the condition on line 198 was never true

199 base = os.environ[envvar] 

200 break 

201 return tempfile.mkdtemp(dir=base) 

202 

203 

204def removeTestTempDir(root: str | None) -> None: 

205 """Attempt to remove a temporary test directory, but do not raise if 

206 unable to. 

207 

208 Unlike `tempfile.TemporaryDirectory`, this passes ``ignore_errors=True`` 

209 to ``shutil.rmtree`` at close, making it safe to use on NFS. 

210 

211 Parameters 

212 ---------- 

213 root : `str`, optional 

214 Name of the directory to be removed. If `None`, nothing will be done. 

215 """ 

216 if root is not None and os.path.exists(root): 216 ↛ exitline 216 didn't return from function 'removeTestTempDir' because the condition on line 216 was always true

217 shutil.rmtree(root, ignore_errors=True) 

218 

219 

220def ensure_directory_is_writeable(directory_path: str | bytes) -> None: 

221 """Given the path to a directory, ensures that we are able to write it and 

222 access files in it. 

223 

224 Alters the directory permissions by adding the owner-write and 

225 owner-traverse permission bits if they aren't already set 

226 

227 Parameters 

228 ---------- 

229 directory_path : `str` or `bytes` 

230 Path to the directory that will be made writeable. 

231 """ 

232 current_mode = os.stat(directory_path).st_mode 

233 desired_mode = current_mode | stat.S_IWUSR | stat.S_IXUSR 

234 if current_mode != desired_mode: 

235 os.chmod(directory_path, desired_mode) 

236 

237 

238def _get_int_env_var(env_var: str) -> int | None: 

239 int_value = None 

240 env_value = os.getenv(env_var) 

241 if env_value is not None: 

242 with contextlib.suppress(TypeError): 

243 int_value = int(env_value) 

244 return int_value 

245 

246 

247@cache 

248def _get_configured_num_workers() -> int | None: 

249 """Return the explicitly requested number of workers. 

250 

251 Returns 

252 ------- 

253 num : `int` or `None` 

254 Value of the ``LSST_RESOURCES_NUM_WORKERS`` environment variable, or 

255 `None` if it is unset or unparsable. 

256 """ 

257 return _get_int_env_var("LSST_RESOURCES_NUM_WORKERS") 

258 

259 

260@cache 

261def _get_default_num_workers() -> int: 

262 """Return the number of workers implied by the available CPUs. 

263 

264 Returns 

265 ------- 

266 num : `int` 

267 The CPU count plus two. Uncapped. 

268 """ 

269 # CPU_LIMIT is used on nublado. 

270 cpu_limit = _get_int_env_var("CPU_LIMIT") or multiprocessing.cpu_count() 

271 return cpu_limit + 2 

272 

273 

274def _get_num_workers(max_workers: int = MAX_WORKERS) -> int: 

275 """Calculate the number of workers to use. 

276 

277 Parameters 

278 ---------- 

279 max_workers : `int`, optional 

280 Upper bound to apply to the calculated default. Ignored when the 

281 number of workers has been requested explicitly. 

282 

283 Returns 

284 ------- 

285 num : `int` 

286 The number of workers to use. The value of 

287 ``$LSST_RESOURCES_NUM_WORKERS`` is used if set, and the CPU count plus 

288 two bounded by ``max_workers`` if not. 

289 """ 

290 configured = _get_configured_num_workers() 

291 if configured is not None: 

292 # An explicit request is honored without capping. 

293 return configured 

294 return min(_get_default_num_workers(), max_workers)