Coverage for src/bioimageio/spec/_internal/io_utils.py: 76%
190 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 09:17 +0000
1from __future__ import annotations
3import collections.abc
4import io
5import shutil
6import zipfile
7from contextlib import nullcontext
8from difflib import get_close_matches
9from pathlib import Path
10from types import MappingProxyType
11from typing import IO, Any, Dict, Mapping, cast
12from zipfile import ZipFile
14import httpx
15import numpy
16from loguru import logger
17from numpy.typing import NDArray
18from pydantic import BaseModel, FilePath, NewPath, RootModel
19from ruyaml import YAML
20from typing_extensions import Unpack
22from ._settings import settings
23from .io import (
24 BIOIMAGEIO_YAML,
25 BioimageioYamlContent,
26 BioimageioYamlContentView,
27 BytesReader,
28 FileDescr,
29 HashKwargs,
30 LightHttpFileDescr,
31 OpenedBioimageioYaml,
32 RelativeFilePath,
33 YamlValue,
34 extract_file_name,
35 find_bioimageio_yaml_file_name,
36 get_reader,
37 identify_bioimageio_yaml_file_name,
38 interprete_file_source,
39)
40from .io_basics import AbsoluteDirectory, FileName, ZipPath
41from .types import PermissiveFileSource
42from .url import HttpUrl, RootHttpUrl
43from .utils import cache
44from .validation_context import ValidationContext, get_validation_context
46_yaml_load = YAML(typ="safe")
48_yaml_dump = YAML()
49_yaml_dump.version = (1, 2) # pyright: ignore[reportAttributeAccessIssue]
50_yaml_dump.default_flow_style = False
51_yaml_dump.indent(mapping=2, sequence=4, offset=2)
52_yaml_dump.width = 88 # pyright: ignore[reportAttributeAccessIssue]
55def read_yaml(
56 file: FilePath | ZipPath | IO[str] | IO[bytes] | BytesReader | str,
57) -> YamlValue:
58 if isinstance(file, (ZipPath, Path)):
59 data = file.read_text(encoding="utf-8")
60 else:
61 data = file
63 content: YamlValue = _yaml_load.load(data)
64 return content
67def write_yaml(
68 content: YamlValue | BioimageioYamlContentView | BaseModel,
69 /,
70 file: NewPath | FilePath | IO[str] | IO[bytes] | ZipPath,
71):
72 if isinstance(file, Path):
73 cm = file.open("w", encoding="utf-8")
74 else:
75 cm = nullcontext(file)
77 if isinstance(content, BaseModel):
78 content = content.model_dump(mode="json")
80 with cm as f:
81 _yaml_dump.dump(content, f)
84def _sanitize_bioimageio_yaml(content: YamlValue) -> BioimageioYamlContent:
85 if not isinstance(content, dict):
86 raise ValueError(
87 f"Expected {BIOIMAGEIO_YAML} content to be a mapping (got {type(content)})."
88 )
90 for key in content:
91 if not isinstance(key, str):
92 raise ValueError(
93 f"Expected all keys (field names) in a {BIOIMAGEIO_YAML} "
94 + f"to be strings (got '{key}' of type {type(key)})."
95 )
97 return cast(BioimageioYamlContent, content)
100def _open_bioimageio_rdf_in_zip(
101 path: ZipPath,
102 *,
103 original_root: AbsoluteDirectory | RootHttpUrl | ZipFile,
104 original_source_name: str,
105) -> OpenedBioimageioYaml:
106 with path.open("rb") as f:
107 assert not isinstance(f, io.TextIOWrapper)
108 unparsed_content = f.read().decode(encoding="utf-8")
110 content = _sanitize_bioimageio_yaml(read_yaml(io.StringIO(unparsed_content)))
112 return OpenedBioimageioYaml(
113 content,
114 original_root=original_root,
115 original_file_name=extract_file_name(path),
116 original_source_name=original_source_name,
117 unparsed_content=unparsed_content,
118 )
121def _open_bioimageio_zip(
122 source: ZipFile,
123 *,
124 original_source_name: str,
125) -> OpenedBioimageioYaml:
126 rdf_name = identify_bioimageio_yaml_file_name(
127 [info.filename for info in source.filelist]
128 )
129 return _open_bioimageio_rdf_in_zip(
130 ZipPath(source, rdf_name),
131 original_root=source,
132 original_source_name=original_source_name,
133 )
136def open_bioimageio_yaml(
137 source: PermissiveFileSource | ZipFile | ZipPath,
138 /,
139 **kwargs: Unpack[HashKwargs],
140) -> OpenedBioimageioYaml:
141 if (
142 isinstance(source, str)
143 and source.startswith("huggingface/")
144 and source.count("/") >= 2
145 ):
146 if source.count("/") == 2:
147 # huggingface/{user_or_org}/{repo_name}
148 repo_id = source[len("huggingface/") :]
149 branch = "main"
150 else:
151 # huggingface/{user_or_org}/{repo_id}/
152 # huggingface/{user_or_org}/{repo_id}/version
153 repo_id, version = source[len("huggingface/") :].rsplit("/", 1)
154 if len(version) == 0:
155 branch = "main"
156 elif version[0].isdigit():
157 branch = f"v{version}"
158 else:
159 branch = version
161 source = HttpUrl(
162 settings.huggingface_http_pattern.format(repo_id=repo_id, branch=branch)
163 )
165 if isinstance(source, RelativeFilePath):
166 source = source.absolute()
168 if isinstance(source, ZipFile):
169 return _open_bioimageio_zip(source, original_source_name=str(source))
170 elif isinstance(source, ZipPath):
171 return _open_bioimageio_rdf_in_zip(
172 source, original_root=source.root, original_source_name=str(source)
173 )
175 try:
176 if isinstance(source, (FileDescr, ZipPath)):
177 src = source
178 elif isinstance(source, (Path, str)) and (source_dir := Path(source)).is_dir():
179 # open bioimageio yaml from a folder
180 src = source_dir / find_bioimageio_yaml_file_name(source_dir)
181 else:
182 src = interprete_file_source(source)
184 reader = get_reader(src, **kwargs)
186 except Exception as e:
187 # check if `source` is a collection id
188 if not isinstance(source, str):
189 raise
191 if settings.collection_http_pattern:
192 with ValidationContext(perform_io_checks=False):
193 url = HttpUrl(
194 settings.collection_http_pattern.format(bioimageio_id=source)
195 )
197 try:
198 r = httpx.get(url, follow_redirects=True)
199 _ = r.raise_for_status()
200 unparsed_content = r.content.decode(encoding="utf-8")
201 content = _sanitize_bioimageio_yaml(read_yaml(unparsed_content))
202 except Exception as e_coll_pattern:
203 collection_pattern_error_msg = f"BIOIMAGEIO_COLLECTION_HTTP_PATTERN: Failed to get bioimageio.yaml from {url}: {e_coll_pattern}"
204 logger.warning(collection_pattern_error_msg)
205 collection_pattern_error_msg = "\n" + collection_pattern_error_msg
206 else:
207 logger.info("loaded {} from {}", source, url)
208 original_file_name = (
209 "rdf.yaml" if url.path is None else url.path.split("/")[-1]
210 )
211 return OpenedBioimageioYaml(
212 content=content,
213 original_root=url.parent,
214 original_file_name=original_file_name,
215 original_source_name=url,
216 unparsed_content=unparsed_content,
217 )
218 else:
219 collection_pattern_error_msg = ""
221 if not isinstance(settings.id_map, str) or "/" not in settings.id_map:
222 raise ValueError(
223 f"BIOIMAGEIO_ID_MAP: Invalid id map url {settings.id_map}.{collection_pattern_error_msg}"
224 ) from e
226 id_map = get_id_map()
227 if not id_map:
228 raise ValueError(
229 f"BIOIMAGEIO_ID_MAP: Empty (or unavailable) id map from {settings.id_map}.{collection_pattern_error_msg}"
230 ) from e
232 if id_map and source not in id_map:
233 close_matches = get_close_matches(source, id_map)
234 if len(close_matches) == 0:
235 raise ValueError(
236 f"BIOIMAGEIO_ID_MAP: '{source}' not found in {settings.id_map}.{collection_pattern_error_msg}"
237 ) from e
239 if len(close_matches) == 1:
240 did_you_mean = f" Did you mean '{close_matches[0]}'?"
241 else:
242 did_you_mean = f" Did you mean any of {close_matches}?"
244 raise ValueError(
245 f"BIOIMAGEIO_ID_MAP: '{source}' not found in {settings.id_map}.{did_you_mean}{collection_pattern_error_msg}"
246 ) from e
248 entry = id_map[source]
249 logger.info("loading {} from {}", source, entry.source)
250 reader = entry.get_reader()
251 with get_validation_context().replace(perform_io_checks=False):
252 src = HttpUrl(entry.source)
254 if reader.is_zipfile:
255 return _open_bioimageio_zip(ZipFile(reader), original_source_name=str(src))
257 unparsed_content = reader.read().decode(encoding="utf-8")
258 content = _sanitize_bioimageio_yaml(read_yaml(unparsed_content))
260 if isinstance(src, RelativeFilePath):
261 src = src.absolute()
263 if isinstance(src, ZipPath):
264 root = src.root
265 elif isinstance(src, FileDescr):
266 file_source = src.source.absolute()
267 if isinstance(file_source, ZipPath):
268 root = file_source.root
269 else:
270 root = file_source.parent
271 else:
272 root = src.parent
274 return OpenedBioimageioYaml(
275 content,
276 original_root=root,
277 original_source_name=str(src),
278 original_file_name=extract_file_name(src),
279 unparsed_content=unparsed_content,
280 )
283_IdMap = RootModel[Dict[str, LightHttpFileDescr]]
286def _get_id_map_impl(url: str) -> dict[str, LightHttpFileDescr]:
287 if not isinstance(url, str) or "/" not in url:
288 logger.opt(depth=1).error("invalid id map url: {}", url)
289 try:
290 id_map_raw: Any = httpx.get(
291 url, timeout=settings.http_timeout, follow_redirects=True
292 ).json()
293 except Exception as e:
294 logger.opt(depth=1).error("failed to get {}: {}", url, e)
295 return {}
297 id_map = _IdMap.model_validate(id_map_raw)
298 return id_map.root
301@cache
302def get_id_map() -> Mapping[str, LightHttpFileDescr]:
303 try:
304 if settings.resolve_draft:
305 ret = _get_id_map_impl(settings.id_map_draft)
306 else:
307 ret = {}
309 ret.update(_get_id_map_impl(settings.id_map))
311 except Exception as e:
312 logger.error("failed to get resource id map: {}", e)
313 ret = {}
315 return MappingProxyType(ret)
318def write_content_to_zip(
319 content: Mapping[
320 FileName,
321 str | FilePath | ZipPath | BioimageioYamlContentView | FileDescr | BytesReader,
322 ],
323 zip: zipfile.ZipFile,
324):
325 """write strings as text, dictionaries as yaml and files to a ZipFile
326 Args:
327 content: dict mapping archive names to local file paths,
328 strings (for text files), or dict (for yaml files).
329 zip: ZipFile
330 """
331 for arc_name, file in content.items():
332 if isinstance(file, collections.abc.Mapping):
333 buf = io.StringIO()
334 write_yaml(file, buf)
335 file = buf.getvalue()
337 if isinstance(file, str):
338 zip.writestr(arc_name, file.encode("utf-8"))
339 else:
340 if isinstance(file, BytesReader):
341 reader = file
342 else:
343 reader = get_reader(file)
345 if (
346 isinstance(reader.original_root, ZipFile)
347 and reader.original_root is zip
348 ):
349 logger.debug(
350 f"Not copying {reader.original_file_name} in "
351 + (
352 "zip file"
353 if reader.original_root.filename is None
354 else reader.original_root.filename
355 )
356 + " to itself."
357 )
358 continue
360 with zip.open(arc_name, "w") as dest:
361 shutil.copyfileobj(reader, dest, 1024 * 8)
364def write_zip(
365 path: FilePath | IO[bytes],
366 content: Mapping[
367 FileName, str | FilePath | ZipPath | BioimageioYamlContentView | BytesReader
368 ],
369 *,
370 compression: int,
371 compression_level: int,
372) -> None:
373 """Write a zip archive.
375 Args:
376 path: output path to write to.
377 content: dict mapping archive names to local file paths, strings (for text files), or dict (for yaml files).
378 compression: The numeric constant of compression method.
379 compression_level: Compression level to use when writing files to the archive.
380 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
382 """
383 if isinstance(path, Path):
384 path.parent.mkdir(parents=True, exist_ok=True)
386 with ZipFile(
387 path, "w", compression=compression, compresslevel=compression_level
388 ) as zip:
389 write_content_to_zip(content, zip)
392def load_array(source: PermissiveFileSource) -> NDArray[Any]:
393 """load a numpy ndarray from a .npy file"""
394 reader = get_reader(source)
395 if settings.allow_pickle:
396 logger.warning("Loading numpy array with `allow_pickle=True`.")
398 return numpy.load(reader, allow_pickle=settings.allow_pickle)
401def save_array(path: Path | ZipPath, array: NDArray[Any]) -> None:
402 """save a numpy ndarray to a .npy file"""
403 with path.open(mode="wb") as f:
404 assert not isinstance(f, io.TextIOWrapper)
405 return numpy.save(f, array, allow_pickle=False)