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