Coverage for src/bioimageio/spec/_internal/utils.py: 65%

1from __future__ import annotations 

2 

3import dataclasses 

4import re 

5import sys 

6from dataclasses import dataclass 

7from functools import wraps 

8from inspect import isfunction, signature 

9from typing import ( 

10 Any, 

11 Callable, 

12 Dict, 

13 Iterable, 

14 TypeVar, 

15) 

16 

17import pydantic 

18from exceptiongroup import ExceptionGroup 

19from ruyaml import Optional 

20from typing_extensions import ParamSpec 

21 

22if sys.version_info < (3, 10): # pragma: no cover 

23 SLOTS: dict[str, bool] = {} 

24else: 

25 SLOTS = {"slots": True} 

26 

27 

28K = TypeVar("K") 

29V = TypeVar("V") 

30NestedDict = Dict[K, "NestedDict[K, V] | V"] 

31 

32if sys.version_info < (3, 9): # pragma: no cover 

33 from functools import lru_cache as cache 

34 from pathlib import Path 

35 

36 def files(package_name: str): 

37 assert package_name == "bioimageio.spec", package_name 

38 return Path(__file__).parent.parent 

39 

40else: 

41 from functools import cache as cache 

42 from importlib.resources import files as files 

43 

44 

45def get_format_version_tuple(format_version: Any) -> Optional[tuple[int, int, int]]: 

46 if ( 

47 not isinstance(format_version, str) 

48 or format_version.count(".") != 2 

49 or any(not v.isdigit() for v in format_version.split(".")) 

50 ): 

51 return None 

52 

53 parsed = tuple(map(int, format_version.split("."))) 

54 assert len(parsed) == 3 

55 return parsed 

56 

57 

58def nest_dict(flat_dict: dict[tuple[K, ...], V]) -> NestedDict[K, V]: 

59 res: NestedDict[K, V] = {} 

60 for k, v in flat_dict.items(): 

61 node: dict[K, NestedDict[K, V] | V] | NestedDict[K, V] = res 

62 for kk in k[:-1]: 

63 if not isinstance(node, dict): 

64 raise ValueError(f"nesting level collision for flat key {k} at {kk}") 

65 d: NestedDict[K, V] = {} 

66 node = node.setdefault(kk, d) # type: ignore 

67 

68 if not isinstance(node, dict): 

69 raise ValueError(f"nesting level collision for flat key {k}") 

70 

71 node[k[-1]] = v 

72 

73 return res 

74 

75 

76FirstK = TypeVar("FirstK") 

77 

78 

79def nest_dict_with_narrow_first_key( 

80 flat_dict: dict[tuple[K, ...], V], first_k: type[FirstK] 

81) -> dict[FirstK, NestedDict[K, V] | V]: 

82 """convenience function to annotate a special version of a NestedDict. 

83 Root level keys are of a narrower type than the nested keys. If not a ValueError is raisd. 

84 """ 

85 nested = nest_dict(flat_dict) 

86 invalid_first_keys = [k for k in nested if not isinstance(k, first_k)] 

87 if invalid_first_keys: 

88 raise ValueError(f"Invalid root level keys: {invalid_first_keys}") 

89 

90 return nested # type: ignore 

91 

92 

93def unindent(text: str, ignore_first_line: bool = False): 

94 """remove minimum count of spaces at beginning of each line. 

95 

96 Args: 

97 text: indented text 

98 ignore_first_line: allows to correctly unindent doc strings 

99 """ 

100 first = int(ignore_first_line) 

101 lines = text.split("\n") 

102 filled_lines = [line for line in lines[first:] if line] 

103 if len(filled_lines) < 2: 

104 return "\n".join(line.strip() for line in lines) 

105 

106 indent = min(len(line) - len(line.lstrip(" ")) for line in filled_lines) 

107 return "\n".join(lines[:first] + [line[indent:] for line in lines[first:]]) 

108 

109 

110T = TypeVar("T") 

111P = ParamSpec("P") 

112 

113 

114def assert_all_params_set_explicitly(fn: Callable[P, T]) -> Callable[P, T]: 

115 @wraps(fn) 

116 def wrapper(*args: P.args, **kwargs: P.kwargs): 

117 n_args = len(args) 

118 missing: set[str] = set() 

119 

120 for p in signature(fn).parameters.values(): 

121 if p.kind == p.POSITIONAL_ONLY: 

122 if n_args == 0: 

123 missing.add(p.name) 

124 else: 

125 n_args -= 1 # 'use' positional arg 

126 elif p.kind == p.POSITIONAL_OR_KEYWORD: 

127 if n_args == 0: 

128 if p.name not in kwargs: 

129 missing.add(p.name) 

130 else: 

131 n_args -= 1 # 'use' positional arg 

132 elif p.kind in (p.VAR_POSITIONAL, p.VAR_KEYWORD): 

133 pass 

134 elif p.kind == p.KEYWORD_ONLY and p.name not in kwargs: 

135 missing.add(p.name) 

136 

137 assert not missing, f"parameters {missing} of {fn} are not set explicitly" 

138 

139 return fn(*args, **kwargs) 

140 

141 return wrapper 

142 

143 

144def get_os_friendly_file_name(name: str) -> str: 

145 return re.sub(r"\W+|^(?=\d)", "_", name) 

146 

147 

148@dataclass 

149class _PrettyDataClassReprMixin: 

150 """A mixin that provides a pretty __repr__ for dataclasses 

151 

152 - leaving out fields that are None 

153 - leaving out memory locations of functions 

154 """ 

155 

156 def __repr__(self): 

157 field_values = { 

158 f.name: v 

159 for f in dataclasses.fields(self) 

160 if (v := getattr(self, f.name)) is not None 

161 } 

162 field_str = ", ".join( 

163 f"{k}=" + (f"<function {v.__name__}>" if isfunction(v) else repr(v)) 

164 for k, v in field_values.items() 

165 ) 

166 return f"{self.__class__.__name__}({field_str})" 

167 

168 

169class PrettyPlainSerializer(pydantic.PlainSerializer, _PrettyDataClassReprMixin): 

170 pass 

171 

172 

173class PrettyWrapSerializer(pydantic.WrapSerializer, _PrettyDataClassReprMixin): 

174 pass 

175 

176 

177def try_all( 

178 funcs: Iterable[Callable[P, T]], 

179 *args: P.args, 

180 **kwargs: P.kwargs, 

181) -> T: 

182 ret, errors = _try_all(funcs, False, *args, **kwargs) 

183 if errors: 

184 raise ExceptionGroup("All functions raised", errors) 

185 

186 assert not isinstance(ret, _AllFailedSentinel) 

187 return ret 

188 

189 

190def try_all_raise_last( 

191 funcs: Iterable[Callable[P, T]], 

192 *args: P.args, 

193 **kwargs: P.kwargs, 

194) -> T: 

195 ret, errors = _try_all(funcs, True, *args, **kwargs) 

196 if errors: 

197 raise errors[-1] 

198 

199 assert not isinstance(ret, _AllFailedSentinel) 

200 return ret 

201 

202 

203class _AllFailedSentinel: 

204 pass 

205 

206 

207def _try_all( 

208 funcs: Iterable[Callable[P, T]], 

209 raise_last_only: bool, 

210 *args: P.args, 

211 **kwargs: P.kwargs, 

212) -> tuple[_AllFailedSentinel | T, list[Exception]]: 

213 """Try to call each of the functions `funcs` with the given arguments. 

214 

215 If all raise, raise an exception group (or the last). 

216 

217 Returns: 

218 Result of the first successful call. 

219 """ 

220 errors: list[Exception] = [] 

221 for c in funcs: 

222 try: 

223 return c(*args, **kwargs), [] 

224 except Exception as e: 

225 errors.append(e) 

226 

227 if errors: 

228 errors.append(RuntimeError("No functions provided to try.")) 

229 

230 return _AllFailedSentinel(), errors