Coverage for src/bioimageio/spec/_internal/io.py: 78%
490 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 hashlib
5import sys
6import warnings
7import zipfile
8from abc import abstractmethod
9from contextlib import nullcontext
10from copy import deepcopy
11from dataclasses import dataclass, field
12from datetime import date as _date
13from datetime import datetime as _datetime
14from functools import partial
15from io import TextIOWrapper
16from pathlib import Path, PurePath, PurePosixPath
17from tempfile import mkdtemp
18from typing import (
19 TYPE_CHECKING,
20 Any,
21 Callable,
22 Dict,
23 Generic,
24 Iterable,
25 List,
26 Mapping,
27 Sequence,
28 Tuple,
29 TypedDict,
30 TypeVar,
31 Union,
32 overload,
33)
34from urllib.parse import urlparse, urlsplit, urlunsplit
35from zipfile import ZipFile
37import httpx
38import pydantic
39from genericache import NoopCache
40from genericache.digest import ContentDigest, UrlDigest
41from pydantic import (
42 AnyUrl,
43 DirectoryPath,
44 Field,
45 GetCoreSchemaHandler,
46 PrivateAttr,
47 RootModel,
48 TypeAdapter,
49 model_serializer,
50 model_validator,
51)
52from pydantic_core import core_schema
53from tqdm import tqdm
54from typing_extensions import (
55 Annotated,
56 LiteralString,
57 NotRequired,
58 Self,
59 TypeAlias,
60 TypeGuard,
61 Unpack,
62 assert_never,
63)
64from typing_extensions import TypeAliasType as _TypeAliasType
66from ._settings import settings
67from .io_basics import (
68 ALL_BIOIMAGEIO_YAML_NAMES,
69 ALTERNATIVE_BIOIMAGEIO_YAML_NAMES,
70 BIOIMAGEIO_YAML,
71 AbsoluteDirectory,
72 AbsoluteFilePath,
73 BytesReader,
74 FileName,
75 FilePath,
76 Sha256,
77 ZipPath,
78 get_sha256,
79)
80from .node import Node
81from .progress import ProgressbarLike
82from .root_url import RootHttpUrl
83from .type_guards import is_dict, is_list, is_mapping, is_sequence
84from .url import HttpUrl
85from .utils import SLOTS
86from .validation_context import get_validation_context
87from .version_type import Version
89AbsolutePathT = TypeVar(
90 "AbsolutePathT",
91 bound=Union[HttpUrl, AbsoluteDirectory, AbsoluteFilePath, ZipPath],
92)
95class LightHttpFileDescr(Node):
96 """http source with sha256 value (minimal validation)"""
98 source: pydantic.HttpUrl
99 """file source"""
101 sha256: Sha256
102 """SHA256 checksum of the source file"""
104 def get_reader(
105 self,
106 *,
107 progressbar: ProgressbarLike
108 | Callable[[], ProgressbarLike]
109 | bool
110 | None = None,
111 ) -> BytesReader:
112 """open the file source (download if needed)"""
113 return get_reader(self.source, sha256=self.sha256, progressbar=progressbar)
115 download = get_reader
116 """alias for get_reader() method"""
119class RelativePathBase(RootModel[PurePath], Generic[AbsolutePathT], frozen=True):
120 _absolute: AbsolutePathT = PrivateAttr()
122 @property
123 def path(self) -> PurePath:
124 return self.root
126 def absolute( # method not property analog to `pathlib.Path.absolute()`
127 self,
128 ) -> AbsolutePathT:
129 """get the absolute path/url
131 (resolved at time of initialization with the root of the ValidationContext)
132 """
133 return self._absolute
135 def model_post_init(self, __context: Any, /) -> None:
136 """set `_absolute` property with validation context at creation time. @private"""
137 if self.root.is_absolute():
138 raise ValueError(f"{self.root} is an absolute path.")
140 if self.root.parts and self.root.parts[0] in ("http:", "https:"):
141 raise ValueError(f"{self.root} looks like an http url.")
143 self._absolute = ( # pyright: ignore[reportAttributeAccessIssue]
144 self.get_absolute(get_validation_context().root)
145 )
146 super().model_post_init(__context)
148 def __str__(self) -> str:
149 return self.root.as_posix()
151 def __repr__(self) -> str:
152 return f"RelativePath('{self}')"
154 @model_serializer()
155 def format(self) -> str:
156 return str(self)
158 @abstractmethod
159 def get_absolute(
160 self, root: RootHttpUrl | AbsoluteDirectory | pydantic.AnyUrl | ZipFile
161 ) -> AbsolutePathT: ...
163 def _get_absolute_impl(
164 self, root: RootHttpUrl | AbsoluteDirectory | pydantic.AnyUrl | ZipFile
165 ) -> Path | HttpUrl | ZipPath:
166 if isinstance(root, Path):
167 return (root / self.root).absolute()
169 rel_path = self.root.as_posix().strip("/")
170 if isinstance(root, ZipFile):
171 return ZipPath(root, rel_path)
173 parsed = urlsplit(str(root))
174 path = list(parsed.path.strip("/").split("/"))
175 if (
176 parsed.netloc == "zenodo.org"
177 and parsed.path.startswith("/api/records/")
178 and parsed.path.endswith("/content")
179 ):
180 path.insert(-1, rel_path)
181 else:
182 path.append(rel_path)
184 return HttpUrl(
185 urlunsplit(
186 (
187 parsed.scheme,
188 parsed.netloc,
189 "/".join(path),
190 parsed.query,
191 parsed.fragment,
192 )
193 )
194 )
196 @classmethod
197 def _validate(cls, value: PurePath | str):
198 if isinstance(value, str) and value.startswith(("https://", "http://")):
199 raise ValueError(f"{value} looks like a URL, not a relative path")
201 return cls(PurePath(value))
204class RelativeFilePath(
205 RelativePathBase[Union[AbsoluteFilePath, HttpUrl, ZipPath]], frozen=True
206):
207 """A path relative to the `rdf.yaml` file (also if the RDF source is a URL)."""
209 def model_post_init(self, __context: Any, /) -> None:
210 """add validation @private"""
211 if not self.root.parts: # an empty path can only be a directory
212 raise ValueError(f"{self.root} is not a valid file path.")
214 super().model_post_init(__context)
216 def get_absolute(
217 self, root: RootHttpUrl | Path | AnyUrl | ZipFile
218 ) -> AbsoluteFilePath | HttpUrl | ZipPath:
219 absolute = self._get_absolute_impl(root)
220 if (
221 isinstance(absolute, Path)
222 and (context := get_validation_context()).perform_io_checks
223 and str(self.root) not in context.known_files
224 and not absolute.is_file()
225 ):
226 raise ValueError(f"{absolute} does not point to an existing file")
228 return absolute
230 @property
231 def suffix(self):
232 return self.root.suffix
235class RelativeDirectory(
236 RelativePathBase[Union[AbsoluteDirectory, HttpUrl, ZipPath]], frozen=True
237):
238 def get_absolute(
239 self, root: RootHttpUrl | Path | AnyUrl | ZipFile
240 ) -> AbsoluteDirectory | HttpUrl | ZipPath:
241 absolute = self._get_absolute_impl(root)
242 if (
243 isinstance(absolute, Path)
244 and get_validation_context().perform_io_checks
245 and not absolute.is_dir()
246 ):
247 raise ValueError(f"{absolute} does not point to an existing directory")
249 return absolute
252FileSource = Annotated[
253 Union[HttpUrl, RelativeFilePath, FilePath],
254 Field(union_mode="left_to_right"),
255]
258class FileDescr(Node):
259 """A file description"""
261 source: FileSource
262 """File source"""
264 sha256: Sha256 | None = None
265 """SHA256 hash value of the **source** file."""
267 @model_validator(mode="after")
268 def _validate_sha256(self) -> Self:
269 self.validate_sha256()
270 return self
272 def validate_sha256(self, force_recompute: bool = False) -> None:
273 """validate the sha256 hash value of the **source** file"""
274 context = get_validation_context()
275 src_str = str(self.source)
276 if force_recompute:
277 actual_sha = None
278 else:
279 actual_sha = context.known_files.get(src_str)
281 if actual_sha is None:
282 if context.perform_io_checks or force_recompute:
283 reader = get_reader(self.source, sha256=self.sha256)
284 if force_recompute:
285 actual_sha = get_sha256(reader)
286 else:
287 actual_sha = reader.sha256
289 context.known_files[src_str] = actual_sha
290 elif context.known_files and src_str not in context.known_files:
291 # perform_io_checks is False, but known files were given,
292 # so we expect all file references to be in there
293 raise ValueError(f"File {src_str} not found in `known_files`.")
295 if actual_sha is None or self.sha256 == actual_sha:
296 return
297 elif self.sha256 is None or context.update_hashes:
298 self.sha256 = actual_sha
299 elif self.sha256 != actual_sha:
300 raise ValueError(
301 f"Sha256 mismatch for {self.source}. Expected {self.sha256}, got "
302 + f"{actual_sha}. Update expected `sha256` or point to the matching "
303 + "file."
304 )
306 def get_reader(
307 self,
308 *,
309 progressbar: ProgressbarLike
310 | Callable[[], ProgressbarLike]
311 | bool
312 | None = None,
313 ):
314 """open the file source (download if needed)"""
315 return get_reader(self.source, progressbar=progressbar, sha256=self.sha256)
317 def download(
318 self,
319 *,
320 progressbar: ProgressbarLike
321 | Callable[[], ProgressbarLike]
322 | bool
323 | None = None,
324 ):
325 """alias for `.get_reader`"""
326 return get_reader(self.source, progressbar=progressbar, sha256=self.sha256)
328 @property
329 def suffix(self) -> str:
330 return self.source.suffix
333PermissiveFileSource: TypeAlias = Union[
334 FileSource, str, pydantic.HttpUrl, FileDescr, ZipPath
335]
338path_or_url_adapter: TypeAdapter[FilePath | DirectoryPath | HttpUrl] = TypeAdapter(
339 Union[FilePath, DirectoryPath, HttpUrl]
340)
343@dataclass(frozen=True, **SLOTS)
344class WithSuffix:
345 suffix: LiteralString | tuple[LiteralString, ...]
346 case_sensitive: bool
348 def __get_pydantic_core_schema__(
349 self, source: type[Any], handler: GetCoreSchemaHandler
350 ):
351 if not self.suffix:
352 raise ValueError("suffix may not be empty")
354 schema = handler(source)
355 return core_schema.no_info_after_validator_function(
356 self.validate,
357 schema,
358 )
360 def validate(self, value: FileSource | FileDescr) -> FileSource | FileDescr:
361 return validate_suffix(value, self.suffix, case_sensitive=self.case_sensitive)
364def wo_special_file_name(src: F) -> F:
365 if has_valid_bioimageio_yaml_name(src):
366 raise ValueError(
367 f"'{src}' not allowed here as its filename is reserved to identify"
368 + f" '{BIOIMAGEIO_YAML}' (or equivalent) files."
369 )
371 return src
374def has_valid_bioimageio_yaml_name(src: FileSource | FileDescr) -> bool:
375 return is_valid_bioimageio_yaml_name(extract_file_name(src))
378def is_valid_bioimageio_yaml_name(file_name: FileName) -> bool:
379 for bioimageio_name in ALL_BIOIMAGEIO_YAML_NAMES:
380 if file_name == bioimageio_name or file_name.endswith("." + bioimageio_name):
381 return True
383 return False
386def identify_bioimageio_yaml_file_name(file_names: Iterable[FileName]) -> FileName:
387 file_names = sorted(file_names)
388 for bioimageio_name in ALL_BIOIMAGEIO_YAML_NAMES:
389 for file_name in file_names:
390 if file_name == bioimageio_name or file_name.endswith(
391 "." + bioimageio_name
392 ):
393 return file_name
395 raise ValueError(
396 f"No {BIOIMAGEIO_YAML} found in {file_names}. (Looking for '{BIOIMAGEIO_YAML}'"
397 + " or or any of the alterntive file names:"
398 + f" {ALTERNATIVE_BIOIMAGEIO_YAML_NAMES}, or any file with an extension of"
399 + f" those, e.g. 'anything.{BIOIMAGEIO_YAML}')."
400 )
403def find_bioimageio_yaml_file_name(path: Path | ZipFile) -> FileName:
404 if isinstance(path, ZipFile):
405 file_names = path.namelist()
406 elif path.is_file():
407 if not zipfile.is_zipfile(path):
408 return path.name
410 with ZipFile(path, "r") as f:
411 file_names = f.namelist()
412 else:
413 file_names = [p.name for p in path.glob("*")]
415 return identify_bioimageio_yaml_file_name(file_names)
418def ensure_has_valid_bioimageio_yaml_name(src: FileSource) -> FileSource:
419 if not has_valid_bioimageio_yaml_name(src):
420 raise ValueError(
421 f"'{src}' does not have a valid filename to identify"
422 + f" '{BIOIMAGEIO_YAML}' (or equivalent) files."
423 )
425 return src
428def ensure_is_valid_bioimageio_yaml_name(file_name: FileName) -> FileName:
429 if not is_valid_bioimageio_yaml_name(file_name):
430 raise ValueError(
431 f"'{file_name}' is not a valid filename to identify"
432 + f" '{BIOIMAGEIO_YAML}' (or equivalent) files."
433 )
435 return file_name
438# types as loaded from YAML 1.2 (with ruyaml)
439YamlLeafValue: TypeAlias = Union[
440 bool, _date, _datetime, int, float, str, None
441] # note: order relevant for deserializing
442YamlKey: TypeAlias = Union[ # YAML Arrays are cast to tuples if used as key in mappings
443 YamlLeafValue, Tuple[YamlLeafValue, ...] # (nesting is not allowed though)
444]
445if TYPE_CHECKING:
446 YamlValue: TypeAlias = Union[
447 YamlLeafValue, List["YamlValue"], Dict[YamlKey, "YamlValue"]
448 ]
449 YamlValueView: TypeAlias = Union[
450 YamlLeafValue, Sequence["YamlValueView"], Mapping[YamlKey, "YamlValueView"]
451 ]
452else:
453 # for pydantic validation we need to use `TypeAliasType`,
454 # see https://docs.pydantic.dev/latest/concepts/types/#named-recursive-types
455 # however this results in a partially unknown type with the current pyright 1.1.388
456 YamlValue: TypeAlias = _TypeAliasType(
457 "YamlValue",
458 Union[YamlLeafValue, List["YamlValue"], Dict[YamlKey, "YamlValue"]],
459 )
460 YamlValueView: TypeAlias = _TypeAliasType(
461 "YamlValueView",
462 Union[
463 YamlLeafValue,
464 Sequence["YamlValueView"],
465 Mapping[YamlKey, "YamlValueView"],
466 ],
467 )
470BioimageioYamlContent = Dict[str, YamlValue]
471BioimageioYamlContentView = Mapping[str, YamlValueView]
473IncompleteDescrLeaf = Union[Node, YamlValue, PermissiveFileSource, Version]
474"""Leaf value of a partial description"""
476IncompleteDescrInner = Union[
477 IncompleteDescrLeaf,
478 List["IncompleteDescrInner"],
479 Dict[YamlKey, "IncompleteDescrInner"],
480]
481"""An inner node of an incomplete resource description --- YAML values and description nodes mixed."""
483IncompleteDescr = Dict[str, IncompleteDescrInner]
484"""An incomplete resource description --- YAML values and description nodes mixed."""
487IncompleteDescrLeafView = Union[Node, YamlValueView, PermissiveFileSource, Version]
488"""Non-editable leaf value of an incomplete description"""
490IncompleteDescrInnerView = Union[
491 IncompleteDescrLeafView,
492 Sequence["IncompleteDescrInnerView"],
493 Mapping[YamlKey, "IncompleteDescrInnerView"],
494 # Mapping[str, YamlValueView], # not sure why this is explicit Mapping is needed
495]
496"""A inner node of a non-editable incomplete resource description --- YAML value views and Node instances mixed."""
498IncompleteDescrView = Mapping[str, IncompleteDescrInnerView]
499"""A non-editable incomplete resource description --- YAML mappings and Node instances mixed."""
502BioimageioYamlSource = Union[
503 PermissiveFileSource, ZipFile, BioimageioYamlContent, BioimageioYamlContentView
504]
507@overload
508def deepcopy_yaml_value(value: BioimageioYamlContentView) -> BioimageioYamlContent: ...
511@overload
512def deepcopy_yaml_value(value: YamlValueView) -> YamlValue: ...
515def deepcopy_yaml_value(
516 value: BioimageioYamlContentView | YamlValueView,
517) -> BioimageioYamlContent | YamlValue:
518 if isinstance(value, collections.abc.Mapping):
519 return {key: deepcopy_yaml_value(val) for key, val in value.items()}
520 elif isinstance(value, collections.abc.Sequence):
521 return [deepcopy_yaml_value(val) for val in value]
522 else:
523 return value
526def deepcopy_incomplete_descr(data: IncompleteDescrView) -> IncompleteDescr:
527 return {k: _deepcopy_incomplete_descr_impl(v) for k, v in data.items()}
530def _deepcopy_incomplete_descr_impl(
531 data: IncompleteDescrInnerView,
532) -> IncompleteDescrInner:
533 if isinstance(data, Node):
534 return deepcopy(data)
535 elif isinstance(data, str):
536 return data
537 elif isinstance(data, collections.abc.Mapping):
538 return {k: _deepcopy_incomplete_descr_impl(v) for k, v in data.items()}
539 elif isinstance(data, collections.abc.Sequence):
540 return [_deepcopy_incomplete_descr_impl(v) for v in data]
541 elif isinstance(
542 data,
543 (
544 HttpUrl,
545 Path,
546 PurePath,
547 RelativeFilePath,
548 Version,
549 _date,
550 _datetime,
551 bool,
552 float,
553 int,
554 pydantic.HttpUrl,
555 type(None),
556 ZipPath,
557 ),
558 ):
559 return data
560 else:
561 assert_never(data)
564def is_yaml_leaf_value(value: Any) -> TypeGuard[YamlLeafValue]:
565 return isinstance(value, (bool, _date, _datetime, int, float, str, type(None)))
568def is_yaml_list(value: Any) -> TypeGuard[list[YamlValue]]:
569 return is_list(value) and all(is_yaml_value(item) for item in value)
572def is_yaml_sequence(value: Any) -> TypeGuard[list[YamlValueView]]:
573 return is_sequence(value) and all(is_yaml_value(item) for item in value)
576def is_yaml_dict(value: Any) -> TypeGuard[BioimageioYamlContent]:
577 return is_dict(value) and all(
578 isinstance(key, str) and is_yaml_value(val) for key, val in value.items()
579 )
582def is_yaml_mapping(value: Any) -> TypeGuard[BioimageioYamlContentView]:
583 return is_mapping(value) and all(
584 isinstance(key, str) and is_yaml_value_read_only(val)
585 for key, val in value.items()
586 )
589def is_yaml_value(value: Any) -> TypeGuard[YamlValue]:
590 return is_yaml_leaf_value(value) or is_yaml_list(value) or is_yaml_dict(value)
593def is_yaml_value_read_only(value: Any) -> TypeGuard[YamlValueView]:
594 return (
595 is_yaml_leaf_value(value) or is_yaml_sequence(value) or is_yaml_mapping(value)
596 )
599@dataclass(frozen=True, **SLOTS)
600class OpenedBioimageioYaml:
601 content: BioimageioYamlContent = field(repr=False)
602 original_root: AbsoluteDirectory | RootHttpUrl | ZipFile
603 original_source_name: str | None
604 original_file_name: FileName
605 unparsed_content: str = field(repr=False)
608@dataclass(frozen=True, **SLOTS)
609class LocalFile:
610 path: FilePath
611 original_root: AbsoluteDirectory | RootHttpUrl | ZipFile
612 original_file_name: FileName
615@dataclass(frozen=True, **SLOTS)
616class FileInZip:
617 path: ZipPath
618 original_root: RootHttpUrl | ZipFile
619 original_file_name: FileName
622class HashKwargs(TypedDict):
623 sha256: NotRequired[Sha256 | None]
626_file_source_adapter: TypeAdapter[HttpUrl | RelativeFilePath | FilePath] = TypeAdapter(
627 FileSource
628)
631def interprete_file_source(
632 file_source: HttpUrl | RelativeFilePath | Path | str | pydantic.HttpUrl,
633) -> HttpUrl | RelativeFilePath | Path:
634 if isinstance(file_source, Path):
635 if (
636 file_source.is_dir()
637 and not file_source.name.endswith(".zarr")
638 and not any(p.name.endswith(".zarr") for p in file_source.parents)
639 ):
640 raise FileNotFoundError(
641 f"{file_source} is a directory, but expected a file (or a '*.zarr' (sub)directory)."
642 )
643 return file_source
645 if isinstance(file_source, HttpUrl):
646 return file_source
648 if isinstance(file_source, pydantic.AnyUrl):
649 file_source = str(file_source)
651 with get_validation_context().replace(perform_io_checks=False):
652 strict = _file_source_adapter.validate_python(file_source)
653 if isinstance(strict, Path) and strict.is_dir():
654 raise FileNotFoundError(
655 f"{strict} is a directory, but expected a file (or a '*.zarr' (sub)directory)."
656 )
658 return strict
661def extract(
662 source: FilePath | ZipFile | ZipPath,
663 folder: DirectoryPath | None = None,
664 overwrite: bool = False,
665) -> DirectoryPath:
666 extract_member = None
667 if isinstance(source, ZipPath):
668 extract_member = source.at
669 source = source.root
671 if isinstance(source, ZipFile):
672 zip_context = nullcontext(source)
673 if folder is None:
674 if source.filename is None:
675 folder = Path(mkdtemp())
676 else:
677 zip_path = Path(source.filename)
678 folder = zip_path.with_suffix(zip_path.suffix + ".unzip")
679 else:
680 zip_context = ZipFile(source, "r")
681 if folder is None:
682 folder = source.with_suffix(source.suffix + ".unzip")
684 if overwrite and folder.exists():
685 warnings.warn(f"Overwriting existing unzipped archive at {folder}")
687 with zip_context as f:
688 if extract_member is not None:
689 extracted_file_path = folder / extract_member
690 if extracted_file_path.exists() and not overwrite:
691 warnings.warn(f"Found unzipped {extracted_file_path}.")
692 else:
693 _ = f.extract(extract_member, folder)
695 return folder
697 elif overwrite or not folder.exists():
698 f.extractall(folder)
699 return folder
701 found_content = {p.relative_to(folder).as_posix() for p in folder.glob("*")}
702 expected_content = {info.filename for info in f.filelist}
703 if expected_missing := expected_content - found_content:
704 parts = folder.name.split("_")
705 nr, *suffixes = parts[-1].split(".")
706 if nr.isdecimal():
707 nr = str(int(nr) + 1)
708 else:
709 nr = f"1.{nr}"
711 parts[-1] = ".".join([nr, *suffixes])
712 out_path_new = folder.with_name("_".join(parts))
713 warnings.warn(
714 f"Unzipped archive at {folder} is missing expected files"
715 + f" {expected_missing}."
716 + f" Unzipping to {out_path_new} instead to avoid overwriting."
717 )
718 return extract(f, out_path_new, overwrite=overwrite)
719 else:
720 warnings.warn(
721 f"Found unzipped archive with all expected files at {folder}."
722 )
723 return folder
726def get_reader(
727 source: PermissiveFileSource | FileDescr | ZipPath,
728 /,
729 progressbar: ProgressbarLike | Callable[[], ProgressbarLike] | bool | None = None,
730 **kwargs: Unpack[HashKwargs],
731) -> BytesReader:
732 """Open a file `source` (download if needed)"""
733 if isinstance(source, FileDescr):
734 if "sha256" not in kwargs:
735 kwargs["sha256"] = source.sha256
737 source = source.source
738 elif isinstance(source, str):
739 source = interprete_file_source(source)
741 if isinstance(source, RelativeFilePath):
742 source = source.absolute()
743 elif isinstance(source, pydantic.AnyUrl):
744 with get_validation_context().replace(perform_io_checks=False):
745 source = HttpUrl(source)
747 if isinstance(source, HttpUrl):
748 return _open_url(source, progressbar=progressbar, **kwargs)
750 if isinstance(source, ZipPath):
751 if not source.exists():
752 raise FileNotFoundError(source.filename)
754 f = source.open(mode="rb")
755 assert not isinstance(f, TextIOWrapper)
756 root = source.root
757 elif isinstance(source, Path):
758 if source.is_dir():
759 raise FileNotFoundError(f"{source} is a directory, not a file")
761 if not source.exists():
762 raise FileNotFoundError(source)
764 f = source.open("rb")
765 root = source.parent
766 else:
767 assert_never(source)
769 expected_sha = kwargs.get("sha256")
770 if expected_sha is None:
771 sha = None
772 else:
773 sha = get_sha256(f)
774 _ = f.seek(0)
775 if sha != expected_sha:
776 raise ValueError(
777 f"SHA256 mismatch for {source}. Expected {expected_sha}, got {sha}."
778 )
780 return BytesReader(
781 f,
782 sha256=sha,
783 suffix=source.suffix,
784 original_file_name=source.name,
785 original_root=root,
786 is_zipfile=None,
787 )
790download = get_reader
793def _open_url(
794 source: HttpUrl,
795 /,
796 progressbar: ProgressbarLike | Callable[[], ProgressbarLike] | bool | None,
797 **kwargs: Unpack[HashKwargs],
798) -> BytesReader:
799 cache = (
800 NoopCache[RootHttpUrl](url_hasher=UrlDigest.from_str)
801 if get_validation_context().disable_cache
802 else settings.disk_cache
803 )
804 sha = kwargs.get("sha256")
805 force_refetch = True if sha is None else ContentDigest.parse(hexdigest=sha)
806 source_path = PurePosixPath(
807 source.path
808 or sha
809 or hashlib.sha256(str(source).encode(encoding="utf-8")).hexdigest()
810 )
812 reader = cache.fetch(
813 source,
814 fetcher=partial(_fetch_url, progressbar=progressbar),
815 force_refetch=force_refetch,
816 )
817 return BytesReader(
818 reader,
819 suffix=source_path.suffix,
820 sha256=sha,
821 original_file_name=source_path.name,
822 original_root=source.parent,
823 is_zipfile=None,
824 )
827def _fetch_url(
828 source: RootHttpUrl,
829 *,
830 progressbar: ProgressbarLike | Callable[[], ProgressbarLike] | bool | None,
831):
832 if source.scheme not in ("http", "https"):
833 raise NotImplementedError(source.scheme)
835 if progressbar is None:
836 # chose progressbar option from validation context
837 progressbar = get_validation_context().progressbar
839 if progressbar is None:
840 # default to no progressbar in CI environments
841 progressbar = not settings.CI
843 if callable(progressbar):
844 progressbar = progressbar()
846 if isinstance(progressbar, bool) and progressbar:
847 progressbar = tqdm(
848 ncols=79,
849 ascii=bool(sys.platform == "win32"),
850 unit="B",
851 unit_scale=True,
852 leave=True,
853 )
855 if progressbar is not False:
856 progressbar.set_description(f"Downloading {extract_file_name(source)}")
858 headers: dict[str, str] = {}
859 if settings.user_agent is not None:
860 headers["User-Agent"] = settings.user_agent
861 elif settings.CI:
862 headers["User-Agent"] = "ci"
864 r = httpx.get(
865 str(source),
866 follow_redirects=True,
867 headers=headers,
868 timeout=settings.http_timeout,
869 )
870 _ = r.raise_for_status()
872 # set progressbar.total
873 total = r.headers.get("content-length")
874 if total is not None and not isinstance(total, int):
875 try:
876 total = int(total)
877 except Exception:
878 total = None
880 if progressbar is not False:
881 if total is None:
882 progressbar.total = 0
883 else:
884 progressbar.total = total
886 def iter_content():
887 for chunk in r.iter_bytes(chunk_size=4096):
888 yield chunk
889 if progressbar is not False:
890 _ = progressbar.update(len(chunk))
892 # Make sure the progress bar gets filled even if the actual number
893 # is chunks is smaller than expected. This happens when streaming
894 # text files that are compressed by the server when sending (gzip).
895 # Binary files don't experience this.
896 # (adapted from pooch.HttpDownloader)
897 if progressbar is not False:
898 progressbar.reset()
899 if total is not None:
900 _ = progressbar.update(total)
902 progressbar.close()
904 return iter_content()
907def extract_file_name(
908 src: pydantic.HttpUrl
909 | RootHttpUrl
910 | PurePath
911 | RelativeFilePath
912 | ZipPath
913 | FileDescr,
914) -> FileName:
915 if isinstance(src, FileDescr):
916 src = src.source
918 if isinstance(src, ZipPath):
919 return src.name or src.root.filename or "bioimageio.zip"
920 elif isinstance(src, RelativeFilePath):
921 return src.path.name
922 elif isinstance(src, PurePath):
923 return src.name
924 else:
925 url = urlparse(str(src))
926 if (
927 url.scheme == "https"
928 and url.hostname == "zenodo.org"
929 and url.path.startswith("/api/records/")
930 and url.path.endswith("/content")
931 ):
932 return url.path.split("/")[-2]
933 else:
934 return url.path.split("/")[-1]
937def extract_file_descrs(
938 data: IncompleteDescrView,
939) -> list[FileDescr]:
940 collected: list[FileDescr] = []
941 with get_validation_context().replace(perform_io_checks=False, log_warnings=False):
942 _extract_file_descrs_impl(data, collected)
944 return collected
947def _extract_file_descrs_impl(
948 data: IncompleteDescrView | IncompleteDescrInnerView,
949 collected: list[FileDescr],
950) -> None:
951 if isinstance(data, FileDescr):
952 collected.append(data)
953 elif isinstance(data, Node):
954 for _, v in data:
955 _extract_file_descrs_impl(v, collected)
956 elif isinstance(data, collections.abc.Mapping):
957 if "source" in data and "sha256" in data:
958 try:
959 fd = FileDescr.model_validate(
960 {"source": data["source"], "sha256": data["sha256"]}
961 )
962 except Exception:
963 warnings.warn(
964 "Found mapping with 'source' and 'sha256' keys, but could not parse it as a FileDescr. Ignoring `sha256`."
965 )
966 try:
967 fd = FileDescr.model_validate({"source": data["source"]})
968 except Exception:
969 warnings.warn(
970 f"Found mapping with 'source' and `sha256' keys , but could not parse it as a FileDescr, evning when ignoring 'sha256'. Ignoring `source`: {data['source']}."
971 )
972 else:
973 collected.append(fd)
974 else:
975 collected.append(fd)
977 for v in data.values():
978 _extract_file_descrs_impl(v, collected)
979 elif not isinstance(data, (str, Path, RelativeFilePath)) and isinstance(
980 data, collections.abc.Sequence
981 ):
982 for v in data:
983 _extract_file_descrs_impl(v, collected)
986F = TypeVar("F", bound=Union[FileSource, FileDescr])
989def validate_suffix(value: F, suffix: str | Sequence[str], case_sensitive: bool) -> F:
990 """check final suffix"""
991 if isinstance(suffix, str):
992 suffixes = [suffix]
993 else:
994 suffixes = suffix
996 assert len(suffixes) > 0, "no suffix given"
997 assert all(suff.startswith(".") for suff in suffixes), (
998 "expected suffixes to start with '.'"
999 )
1000 o_value = value
1001 if isinstance(value, FileDescr):
1002 strict = value.source
1003 else:
1004 strict = interprete_file_source(value)
1006 if isinstance(strict, (HttpUrl, AnyUrl)):
1007 if strict.path is None or "." not in (path := strict.path):
1008 actual_suffixes = []
1009 else:
1010 if (
1011 strict.host == "zenodo.org"
1012 and path.startswith("/api/records/")
1013 and path.endswith("/content")
1014 ):
1015 # Zenodo API URLs have a "/content" suffix that should be ignored
1016 path = path[: -len("/content")]
1018 actual_suffixes = [f".{path.split('.')[-1]}"]
1020 elif isinstance(strict, PurePath):
1021 actual_suffixes = strict.suffixes
1022 elif isinstance(strict, RelativeFilePath):
1023 actual_suffixes = strict.path.suffixes
1024 else:
1025 assert_never(strict)
1027 if actual_suffixes:
1028 actual_suffix = actual_suffixes[-1]
1029 else:
1030 actual_suffix = "no suffix"
1032 if (
1033 case_sensitive
1034 and actual_suffix not in suffixes
1035 or not case_sensitive
1036 and actual_suffix.lower() not in [s.lower() for s in suffixes]
1037 ):
1038 if len(suffixes) == 1:
1039 raise ValueError(f"Expected suffix {suffixes[0]}, but got {actual_suffix}")
1040 else:
1041 raise ValueError(
1042 f"Expected a suffix from {suffixes}, but got {actual_suffix}"
1043 )
1045 return o_value
1048def populate_cache(sources: Sequence[FileDescr | LightHttpFileDescr]):
1049 unique: set[str] = set()
1050 for src in sources:
1051 if src.sha256 is None:
1052 continue # not caching without known SHA
1054 if isinstance(src.source, (HttpUrl, pydantic.AnyUrl)):
1055 url = str(src.source)
1056 elif isinstance(src.source, RelativeFilePath):
1057 if isinstance(absolute := src.source.absolute(), HttpUrl):
1058 url = str(absolute)
1059 else:
1060 continue # not caching local paths
1061 elif isinstance(src.source, Path):
1062 continue # not caching local paths
1063 else:
1064 assert_never(src.source)
1066 if url in unique:
1067 continue # skip duplicate URLs
1069 unique.add(url)
1070 _ = src.download()