Coverage for src/bioimageio/spec/_internal/io.py: 75%
525 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 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 Literal,
27 Mapping,
28 Sequence,
29 Tuple,
30 TypedDict,
31 TypeVar,
32 Union,
33 overload,
34)
35from urllib.parse import urlparse, urlsplit, urlunsplit
36from zipfile import ZipFile
38import httpx
39import pydantic
40from exceptiongroup import ExceptionGroup
41from genericache import NoopCache
42from genericache.digest import ContentDigest, UrlDigest
43from pydantic import (
44 DirectoryPath,
45 Field,
46 GetCoreSchemaHandler,
47 PrivateAttr,
48 RootModel,
49 TypeAdapter,
50 model_serializer,
51 model_validator,
52)
53from pydantic_core import core_schema
54from tqdm import tqdm
55from typing_extensions import (
56 Annotated,
57 LiteralString,
58 NotRequired,
59 Self,
60 TypeAlias,
61 TypeGuard,
62 Unpack,
63 assert_never,
64)
65from typing_extensions import TypeAliasType as _TypeAliasType
67from ._settings import settings
68from .io_basics import (
69 ALL_BIOIMAGEIO_YAML_NAMES,
70 ALTERNATIVE_BIOIMAGEIO_YAML_NAMES,
71 BIOIMAGEIO_YAML,
72 AbsoluteDirectory,
73 AbsoluteFilePath,
74 BytesReader,
75 FileName,
76 FilePath,
77 Sha256,
78 ZipPath,
79 get_sha256,
80)
81from .node import Node
82from .progress import ProgressbarLike
83from .root_url import FtpUrl, RootHttpUrl
84from .type_guards import is_dict, is_list, is_mapping, is_sequence
85from .url import HttpUrl
86from .utils import SLOTS
87from .validation_context import get_validation_context
88from .version_type import Version
90AbsolutePathT = TypeVar(
91 "AbsolutePathT",
92 bound=Union[HttpUrl, AbsoluteDirectory, AbsoluteFilePath, ZipPath],
93)
96class LightHttpFileDescr(Node):
97 """http source with sha256 value (minimal validation)"""
99 source: pydantic.HttpUrl
100 """file source"""
102 sha256: Sha256
103 """SHA256 checksum of the source file"""
105 def get_reader(
106 self,
107 *,
108 progressbar: ProgressbarLike
109 | Callable[[], ProgressbarLike]
110 | bool
111 | None = None,
112 ) -> BytesReader:
113 """open the file source (download if needed)"""
114 return get_reader(self.source, sha256=self.sha256, progressbar=progressbar)
116 download = get_reader
117 """alias for get_reader() method"""
120class RelativePathBase(RootModel[PurePath], Generic[AbsolutePathT], frozen=True):
121 _absolute: AbsolutePathT = PrivateAttr()
123 @property
124 def path(self) -> PurePath:
125 return self.root
127 @property
128 def suffix(self):
129 return self.root.suffix
131 def absolute( # method not property analog to `pathlib.Path.absolute()`
132 self,
133 ) -> AbsolutePathT:
134 """get the absolute path/url
136 (resolved at time of initialization with the root of the ValidationContext)
137 """
138 return self._absolute
140 def model_post_init(self, __context: Any, /) -> None:
141 """set `_absolute` property with validation context at creation time. @private"""
142 if self.root.is_absolute():
143 raise ValueError(f"{self.root} is an absolute path.")
145 if self.root.parts and self.root.parts[0] in ("http:", "https:"):
146 raise ValueError(f"{self.root} looks like an http url.")
148 self._absolute = ( # pyright: ignore[reportAttributeAccessIssue]
149 self.get_absolute(get_validation_context().root)
150 )
151 super().model_post_init(__context)
153 def __str__(self) -> str:
154 return self.root.as_posix()
156 def __repr__(self) -> str:
157 return f"RelativePath('{self}')"
159 @model_serializer()
160 def format(self) -> str:
161 return str(self)
163 @abstractmethod
164 def get_absolute(
165 self, root: RootHttpUrl | AbsoluteDirectory | pydantic.AnyUrl | ZipFile
166 ) -> AbsolutePathT: ...
168 def _get_absolute_impl(
169 self, root: RootHttpUrl | AbsoluteDirectory | pydantic.AnyUrl | ZipFile
170 ) -> Path | HttpUrl | ZipPath:
171 if isinstance(root, Path):
172 return (root / self.root).absolute()
174 rel_path = self.root.as_posix().strip("/")
175 if isinstance(root, ZipFile):
176 return ZipPath(root, rel_path)
178 parsed = urlsplit(str(root))
179 path = list(parsed.path.strip("/").split("/"))
180 if (
181 parsed.netloc == "zenodo.org"
182 and parsed.path.startswith("/api/records/")
183 and parsed.path.endswith("/content")
184 ):
185 path.insert(-1, rel_path)
186 else:
187 path.append(rel_path)
189 return HttpUrl(
190 urlunsplit(
191 (
192 parsed.scheme,
193 parsed.netloc,
194 "/".join(path),
195 parsed.query,
196 parsed.fragment,
197 )
198 )
199 )
201 @classmethod
202 def _validate(cls, value: PurePath | str):
203 if isinstance(value, str) and value.startswith(("https://", "http://")):
204 raise ValueError(f"{value} looks like a URL, not a relative path")
206 return cls(PurePath(value))
209class RelativeFilePath(
210 RelativePathBase[Union[AbsoluteFilePath, HttpUrl, ZipPath]], frozen=True
211):
212 """A path relative to the `rdf.yaml` file (also if the RDF source is a URL)."""
214 def model_post_init(self, __context: Any, /) -> None:
215 """add validation @private"""
216 if not self.root.parts: # an empty path can only be a directory
217 raise ValueError(f"{self.root} is not a valid file path.")
219 super().model_post_init(__context)
221 def get_absolute(
222 self, root: RootHttpUrl | Path | pydantic.AnyUrl | ZipFile
223 ) -> AbsoluteFilePath | HttpUrl | ZipPath:
224 absolute = self._get_absolute_impl(root)
225 if (
226 isinstance(absolute, Path)
227 and (context := get_validation_context()).perform_io_checks
228 and str(self.root) not in context.known_files
229 and not absolute.is_file()
230 ):
231 raise ValueError(f"{absolute} does not point to an existing file")
233 return absolute
235 @property
236 def parent(self) -> RelativeDirectory:
237 return RelativeDirectory(self.root.parent)
240class RelativeDirectory(
241 RelativePathBase[Union[AbsoluteDirectory, HttpUrl, ZipPath]], frozen=True
242):
243 def get_absolute(
244 self, root: RootHttpUrl | Path | pydantic.AnyUrl | ZipFile
245 ) -> AbsoluteDirectory | HttpUrl | ZipPath:
246 absolute = self._get_absolute_impl(root)
247 if (
248 isinstance(absolute, Path)
249 and get_validation_context().perform_io_checks
250 and not absolute.is_dir()
251 ):
252 raise ValueError(f"{absolute} does not point to an existing directory")
254 return absolute
256 @property
257 def parent(self) -> RelativeDirectory:
258 return RelativeDirectory(self.root.parent)
261@dataclass(frozen=True, **SLOTS)
262class WithSuffix:
263 suffix: LiteralString | tuple[LiteralString, ...]
264 case_sensitive: bool
265 allow_any_parent_suffix: bool = False
266 """Considers the suffix of any parent directory as well, e.g. for `foo.zarr/bar` or `foo.zarr/bar/buz`"""
268 def __get_pydantic_core_schema__(
269 self, source: type[Any], handler: GetCoreSchemaHandler
270 ):
271 if not self.suffix:
272 raise ValueError("suffix may not be empty")
274 schema = handler(source)
275 return core_schema.no_info_after_validator_function(
276 self.validate,
277 schema,
278 )
280 def validate(self, value: FileSource | FileDescr) -> FileSource | FileDescr:
281 return validate_suffix(value, self.suffix, case_sensitive=self.case_sensitive)
284ZarrPath = Annotated[
285 pydantic.DirectoryPath,
286 pydantic.Field(title="ZarrPath"),
287 WithSuffix(".zarr", case_sensitive=True),
288]
289RelativeZarrPath = Annotated[
290 RelativeDirectory,
291 pydantic.Field(title="RelativeZarrPath"),
292 WithSuffix(".zarr", case_sensitive=True),
293]
295ZarrUrl: TypeAlias = Annotated[
296 Union[RootHttpUrl, FtpUrl],
297 pydantic.Field(title="ZarrUrl"),
298 WithSuffix(".zarr", case_sensitive=True, allow_any_parent_suffix=True),
299]
300"""An untested HTTP/FTP URL to a zarr (sub) directory"""
303FileSource: TypeAlias = Annotated[
304 Union[HttpUrl, RelativeFilePath, FilePath],
305 Field(title="FileSource", union_mode="left_to_right"),
306]
307ZarrSource: TypeAlias = Annotated[
308 Union[ZarrUrl, RelativeZarrPath, ZarrPath],
309 Field(title="ZarrSource", union_mode="left_to_right"),
310]
313class FileDescr(Node):
314 """A file description"""
316 source: FileSource
317 """File source"""
319 sha256: Sha256 | None = None
320 """SHA256 hash value of the **source** file."""
322 @model_validator(mode="after")
323 def _validate_sha256(self) -> Self:
324 self.validate_sha256()
325 return self
327 def validate_sha256(self, force_recompute: bool = False) -> None:
328 """validate the sha256 hash value of the **source** file"""
329 context = get_validation_context()
330 src_str = str(self.source)
331 if force_recompute:
332 actual_sha = None
333 else:
334 actual_sha = context.known_files.get(src_str)
336 if actual_sha is None:
337 if context.perform_io_checks or force_recompute:
338 reader = get_reader(self.source, sha256=self.sha256)
339 if force_recompute:
340 actual_sha = get_sha256(reader)
341 else:
342 actual_sha = reader.sha256
344 context.known_files[src_str] = actual_sha
345 elif context.known_files and src_str not in context.known_files:
346 # perform_io_checks is False, but known files were given,
347 # so we expect all file references to be in there
348 raise ValueError(f"File {src_str} not found in `known_files`.")
350 if actual_sha is None or self.sha256 == actual_sha:
351 return
352 elif self.sha256 is None or context.update_hashes:
353 self.sha256 = actual_sha
354 elif self.sha256 != actual_sha:
355 raise ValueError(
356 f"Sha256 mismatch for {self.source}. Expected {self.sha256}, got "
357 + f"{actual_sha}. Update expected `sha256` or point to the matching "
358 + "file."
359 )
361 def get_reader(
362 self,
363 *,
364 progressbar: ProgressbarLike
365 | Callable[[], ProgressbarLike]
366 | bool
367 | None = None,
368 ):
369 """open the file source (download if needed)"""
370 return get_reader(self.source, progressbar=progressbar, sha256=self.sha256)
372 def download(
373 self,
374 *,
375 progressbar: ProgressbarLike
376 | Callable[[], ProgressbarLike]
377 | bool
378 | None = None,
379 ):
380 """alias for `.get_reader`"""
381 return get_reader(self.source, progressbar=progressbar, sha256=self.sha256)
383 @property
384 def suffix(self) -> str:
385 return self.source.suffix
388PermissiveFileSource: TypeAlias = Union[
389 FileSource, str, pydantic.HttpUrl, FileDescr, ZipPath
390]
393path_or_url_adapter: TypeAdapter[FilePath | DirectoryPath | HttpUrl] = TypeAdapter(
394 Union[FilePath, DirectoryPath, HttpUrl]
395)
398def wo_special_file_name(src: F) -> F:
399 if has_valid_bioimageio_yaml_name(src):
400 raise ValueError(
401 f"'{src}' not allowed here as its filename is reserved to identify"
402 + f" '{BIOIMAGEIO_YAML}' (or equivalent) files."
403 )
405 return src
408def has_valid_bioimageio_yaml_name(
409 src: ZarrSource | FileSource | FileDescr,
410) -> bool:
411 return is_valid_bioimageio_yaml_name(extract_file_name(src))
414def is_valid_bioimageio_yaml_name(file_name: FileName) -> bool:
415 for bioimageio_name in ALL_BIOIMAGEIO_YAML_NAMES:
416 if file_name == bioimageio_name or file_name.endswith("." + bioimageio_name):
417 return True
419 return False
422def identify_bioimageio_yaml_file_name(file_names: Iterable[FileName]) -> FileName:
423 file_names = sorted(file_names)
424 for bioimageio_name in ALL_BIOIMAGEIO_YAML_NAMES:
425 for file_name in file_names:
426 if file_name == bioimageio_name or file_name.endswith(
427 "." + bioimageio_name
428 ):
429 return file_name
431 raise ValueError(
432 f"No {BIOIMAGEIO_YAML} found in {file_names}. (Looking for '{BIOIMAGEIO_YAML}'"
433 + " or or any of the alterntive file names:"
434 + f" {ALTERNATIVE_BIOIMAGEIO_YAML_NAMES}, or any file with an extension of"
435 + f" those, e.g. 'anything.{BIOIMAGEIO_YAML}')."
436 )
439def find_bioimageio_yaml_file_name(path: Path | ZipFile) -> FileName:
440 if isinstance(path, ZipFile):
441 file_names = path.namelist()
442 elif path.is_file():
443 if not zipfile.is_zipfile(path):
444 return path.name
446 with ZipFile(path, "r") as f:
447 file_names = f.namelist()
448 else:
449 file_names = [p.name for p in path.glob("*")]
451 return identify_bioimageio_yaml_file_name(file_names)
454def ensure_has_valid_bioimageio_yaml_name(src: FileSource) -> FileSource:
455 if not has_valid_bioimageio_yaml_name(src):
456 raise ValueError(
457 f"'{src}' does not have a valid filename to identify"
458 + f" '{BIOIMAGEIO_YAML}' (or equivalent) files."
459 )
461 return src
464def ensure_is_valid_bioimageio_yaml_name(file_name: FileName) -> FileName:
465 if not is_valid_bioimageio_yaml_name(file_name):
466 raise ValueError(
467 f"'{file_name}' is not a valid filename to identify"
468 + f" '{BIOIMAGEIO_YAML}' (or equivalent) files."
469 )
471 return file_name
474# types as loaded from YAML 1.2 (with ruyaml)
475YamlLeafValue: TypeAlias = Union[
476 bool, _date, _datetime, int, float, str, None
477] # note: order relevant for deserializing
478YamlKey: TypeAlias = Union[ # YAML Arrays are cast to tuples if used as key in mappings
479 YamlLeafValue, Tuple[YamlLeafValue, ...] # (nesting is not allowed though)
480]
481if TYPE_CHECKING:
482 YamlValue: TypeAlias = Union[
483 YamlLeafValue, List["YamlValue"], Dict[YamlKey, "YamlValue"]
484 ]
485 YamlValueView: TypeAlias = Union[
486 YamlLeafValue, Sequence["YamlValueView"], Mapping[YamlKey, "YamlValueView"]
487 ]
488else:
489 # for pydantic validation we need to use `TypeAliasType`,
490 # see https://docs.pydantic.dev/latest/concepts/types/#named-recursive-types
491 # however this results in a partially unknown type with the current pyright 1.1.388
492 YamlValue: TypeAlias = _TypeAliasType(
493 "YamlValue",
494 Union[YamlLeafValue, List["YamlValue"], Dict[YamlKey, "YamlValue"]],
495 )
496 YamlValueView: TypeAlias = _TypeAliasType(
497 "YamlValueView",
498 Union[
499 YamlLeafValue,
500 Sequence["YamlValueView"],
501 Mapping[YamlKey, "YamlValueView"],
502 ],
503 )
506BioimageioYamlContent = Dict[str, YamlValue]
507BioimageioYamlContentView = Mapping[str, YamlValueView]
509IncompleteDescrLeaf = Union[Node, YamlValue, PermissiveFileSource, Version]
510"""Leaf value of a partial description"""
512IncompleteDescrInner = Union[
513 IncompleteDescrLeaf,
514 List["IncompleteDescrInner"],
515 Dict[YamlKey, "IncompleteDescrInner"],
516]
517"""An inner node of an incomplete resource description --- YAML values and description nodes mixed."""
519IncompleteDescr = Dict[str, IncompleteDescrInner]
520"""An incomplete resource description --- YAML values and description nodes mixed."""
523IncompleteDescrLeafView = Union[Node, YamlValueView, PermissiveFileSource, Version]
524"""Non-editable leaf value of an incomplete description"""
526IncompleteDescrInnerView = Union[
527 IncompleteDescrLeafView,
528 Sequence["IncompleteDescrInnerView"],
529 Mapping[YamlKey, "IncompleteDescrInnerView"],
530 # Mapping[str, YamlValueView], # not sure why this is explicit Mapping is needed
531]
532"""A inner node of a non-editable incomplete resource description --- YAML value views and Node instances mixed."""
534IncompleteDescrView = Mapping[str, IncompleteDescrInnerView]
535"""A non-editable incomplete resource description --- YAML mappings and Node instances mixed."""
538BioimageioYamlSource: TypeAlias = Union[
539 PermissiveFileSource, ZipFile, BioimageioYamlContent, BioimageioYamlContentView
540]
543@overload
544def deepcopy_yaml_value(value: BioimageioYamlContentView) -> BioimageioYamlContent: ...
547@overload
548def deepcopy_yaml_value(value: YamlValueView) -> YamlValue: ...
551def deepcopy_yaml_value(
552 value: BioimageioYamlContentView | YamlValueView,
553) -> BioimageioYamlContent | YamlValue:
554 if isinstance(value, collections.abc.Mapping):
555 return {key: deepcopy_yaml_value(val) for key, val in value.items()}
556 elif isinstance(value, collections.abc.Sequence):
557 return [deepcopy_yaml_value(val) for val in value]
558 else:
559 return value
562def deepcopy_incomplete_descr(data: IncompleteDescrView) -> IncompleteDescr:
563 return {k: _deepcopy_incomplete_descr_impl(v) for k, v in data.items()}
566def _deepcopy_incomplete_descr_impl(
567 data: IncompleteDescrInnerView,
568) -> IncompleteDescrInner:
569 if isinstance(data, Node):
570 return deepcopy(data)
571 elif isinstance(data, str):
572 return data
573 elif isinstance(data, collections.abc.Mapping):
574 return {k: _deepcopy_incomplete_descr_impl(v) for k, v in data.items()}
575 elif isinstance(data, collections.abc.Sequence):
576 return [_deepcopy_incomplete_descr_impl(v) for v in data]
577 elif isinstance(
578 data,
579 (
580 HttpUrl,
581 Path,
582 PurePath,
583 RelativeFilePath,
584 RelativeDirectory,
585 Version,
586 _date,
587 _datetime,
588 bool,
589 float,
590 int,
591 pydantic.HttpUrl,
592 type(None),
593 ZipPath,
594 ),
595 ):
596 return data
597 else:
598 assert_never(data)
601def is_yaml_leaf_value(value: Any) -> TypeGuard[YamlLeafValue]:
602 return isinstance(value, (bool, _date, _datetime, int, float, str, type(None)))
605def is_yaml_list(value: Any) -> TypeGuard[list[YamlValue]]:
606 return is_list(value) and all(is_yaml_value(item) for item in value)
609def is_yaml_sequence(value: Any) -> TypeGuard[list[YamlValueView]]:
610 return is_sequence(value) and all(is_yaml_value(item) for item in value)
613def is_yaml_dict(value: Any) -> TypeGuard[BioimageioYamlContent]:
614 return is_dict(value) and all(
615 isinstance(key, str) and is_yaml_value(val) for key, val in value.items()
616 )
619def is_yaml_mapping(value: Any) -> TypeGuard[BioimageioYamlContentView]:
620 return is_mapping(value) and all(
621 isinstance(key, str) and is_yaml_value_read_only(val)
622 for key, val in value.items()
623 )
626def is_yaml_value(value: Any) -> TypeGuard[YamlValue]:
627 return is_yaml_leaf_value(value) or is_yaml_list(value) or is_yaml_dict(value)
630def is_yaml_value_read_only(value: Any) -> TypeGuard[YamlValueView]:
631 return (
632 is_yaml_leaf_value(value) or is_yaml_sequence(value) or is_yaml_mapping(value)
633 )
636@dataclass(frozen=True, **SLOTS)
637class OpenedBioimageioYaml:
638 content: BioimageioYamlContent = field(repr=False)
639 original_root: AbsoluteDirectory | RootHttpUrl | ZipFile
640 original_source_name: str | None
641 original_file_name: FileName
642 unparsed_content: str = field(repr=False)
645@dataclass(frozen=True, **SLOTS)
646class LocalFile:
647 path: FilePath
648 original_root: AbsoluteDirectory | RootHttpUrl | ZipFile
649 original_file_name: FileName
652@dataclass(frozen=True, **SLOTS)
653class FileInZip:
654 path: ZipPath
655 original_root: RootHttpUrl | ZipFile
656 original_file_name: FileName
659class HashKwargs(TypedDict):
660 sha256: NotRequired[Sha256 | None]
663_file_source_adapter: TypeAdapter[FileSource] = TypeAdapter(FileSource)
664_zarr_source_adapter: TypeAdapter[ZarrSource] = TypeAdapter(ZarrSource)
667@overload
668def interprete_file_source(
669 file_source: str | pydantic.AnyUrl, allow_zarr: Literal[False] = False
670) -> FileSource: ...
673@overload
674def interprete_file_source(
675 file_source: str | pydantic.AnyUrl, allow_zarr: Literal[True] = True
676) -> FileSource | ZarrSource: ...
679def interprete_file_source(
680 file_source: str | pydantic.AnyUrl, allow_zarr: bool = False
681) -> FileSource | ZarrSource:
682 file_source = str(file_source)
683 with get_validation_context().replace(perform_io_checks=False):
684 try:
685 strict = _file_source_adapter.validate_python(file_source)
686 except Exception as e1:
687 if allow_zarr:
688 try:
689 strict = _zarr_source_adapter.validate_python(file_source)
690 except Exception as e2:
691 raise ExceptionGroup(
692 f"Could not interpret {file_source} as a file source or zarr source.",
693 (e1, e2),
694 )
695 else:
696 raise
698 return strict
701def interprete_zarr_source(
702 zarr_source: str | pydantic.AnyUrl,
703) -> ZarrSource:
704 zarr_source = str(zarr_source)
705 with get_validation_context().replace(perform_io_checks=False):
706 strict = _zarr_source_adapter.validate_python(zarr_source)
708 return strict
711def extract(
712 source: FilePath | ZipFile | ZipPath,
713 folder: DirectoryPath | None = None,
714 overwrite: bool = False,
715) -> DirectoryPath:
716 extract_member = None
717 if isinstance(source, ZipPath):
718 extract_member = source.at
719 source = source.root
721 if isinstance(source, ZipFile):
722 zip_context = nullcontext(source)
723 if folder is None:
724 if source.filename is None:
725 folder = Path(mkdtemp())
726 else:
727 zip_path = Path(source.filename)
728 folder = zip_path.with_suffix(zip_path.suffix + ".unzip")
729 else:
730 zip_context = ZipFile(source, "r")
731 if folder is None:
732 folder = source.with_suffix(source.suffix + ".unzip")
734 if overwrite and folder.exists():
735 warnings.warn(f"Overwriting existing unzipped archive at {folder}")
737 with zip_context as f:
738 if extract_member is not None:
739 extracted_file_path = folder / extract_member
740 if extracted_file_path.exists() and not overwrite:
741 warnings.warn(f"Found unzipped {extracted_file_path}.")
742 else:
743 _ = f.extract(extract_member, folder)
745 return folder
747 elif overwrite or not folder.exists():
748 f.extractall(folder)
749 return folder
751 found_content = {p.relative_to(folder).as_posix() for p in folder.glob("*")}
752 expected_content = {info.filename for info in f.filelist}
753 if expected_missing := expected_content - found_content:
754 parts = folder.name.split("_")
755 nr, *suffixes = parts[-1].split(".")
756 if nr.isdecimal():
757 nr = str(int(nr) + 1)
758 else:
759 nr = f"1.{nr}"
761 parts[-1] = ".".join([nr, *suffixes])
762 out_path_new = folder.with_name("_".join(parts))
763 warnings.warn(
764 f"Unzipped archive at {folder} is missing expected files"
765 + f" {expected_missing}."
766 + f" Unzipping to {out_path_new} instead to avoid overwriting."
767 )
768 return extract(f, out_path_new, overwrite=overwrite)
769 else:
770 warnings.warn(
771 f"Found unzipped archive with all expected files at {folder}."
772 )
773 return folder
776def get_reader(
777 source: PermissiveFileSource | FileDescr | ZipPath,
778 /,
779 progressbar: ProgressbarLike | Callable[[], ProgressbarLike] | bool | None = None,
780 **kwargs: Unpack[HashKwargs],
781) -> BytesReader:
782 """Open a file `source` (download if needed)"""
783 if isinstance(source, FileDescr):
784 if "sha256" not in kwargs:
785 kwargs["sha256"] = source.sha256
787 source = source.source
788 elif isinstance(source, str):
789 source = interprete_file_source(source)
791 if isinstance(source, RelativeDirectory):
792 raise ValueError(f"{source} is not a file source, but a directory")
793 elif isinstance(source, RelativeFilePath):
794 source = source.absolute()
795 elif isinstance(source, (pydantic.AnyUrl, RootHttpUrl)):
796 with get_validation_context().replace(perform_io_checks=False):
797 source = HttpUrl(source)
799 if isinstance(source, FtpUrl):
800 raise NotImplementedError(
801 "FTP URLs are not supported yet. Please use HTTP(S) URLs instead."
802 )
804 if isinstance(source, HttpUrl):
805 return _open_url(source, progressbar=progressbar, **kwargs)
807 if isinstance(source, ZipPath):
808 if not source.exists():
809 raise FileNotFoundError(source.filename)
811 f = source.open(mode="rb")
812 assert not isinstance(f, TextIOWrapper)
813 root = source.root
814 elif isinstance(source, Path):
815 if source.is_dir():
816 raise FileNotFoundError(f"{source} is a directory, not a file")
818 if not source.exists():
819 raise FileNotFoundError(source)
821 f = source.open("rb")
822 root = source.parent
823 else:
824 assert_never(source)
826 expected_sha = kwargs.get("sha256")
827 if expected_sha is None:
828 sha = None
829 else:
830 sha = get_sha256(f)
831 _ = f.seek(0)
832 if sha != expected_sha:
833 raise ValueError(
834 f"SHA256 mismatch for {source}. Expected {expected_sha}, got {sha}."
835 )
837 return BytesReader(
838 f,
839 sha256=sha,
840 suffix=source.suffix,
841 original_file_name=source.name,
842 original_root=root,
843 is_zipfile=None,
844 )
847download = get_reader
850def _open_url(
851 source: HttpUrl,
852 /,
853 progressbar: ProgressbarLike | Callable[[], ProgressbarLike] | bool | None,
854 **kwargs: Unpack[HashKwargs],
855) -> BytesReader:
856 cache = (
857 NoopCache[RootHttpUrl](url_hasher=UrlDigest.from_str)
858 if get_validation_context().disable_cache
859 else settings.disk_cache
860 )
861 sha = kwargs.get("sha256")
862 force_refetch = True if sha is None else ContentDigest.parse(hexdigest=sha)
863 source_path = PurePosixPath(
864 source.path
865 or sha
866 or hashlib.sha256(str(source).encode(encoding="utf-8")).hexdigest()
867 )
869 reader = cache.fetch(
870 source,
871 fetcher=partial(_fetch_url, progressbar=progressbar),
872 force_refetch=force_refetch,
873 )
874 return BytesReader(
875 reader,
876 suffix=source_path.suffix,
877 sha256=sha,
878 original_file_name=source_path.name,
879 original_root=source.parent,
880 is_zipfile=None,
881 )
884def _fetch_url(
885 source: RootHttpUrl,
886 *,
887 progressbar: ProgressbarLike | Callable[[], ProgressbarLike] | bool | None,
888):
889 if source.scheme not in ("http", "https"):
890 raise NotImplementedError(source.scheme)
892 if progressbar is None:
893 # chose progressbar option from validation context
894 progressbar = get_validation_context().progressbar
896 if progressbar is None:
897 # default to no progressbar in CI environments
898 progressbar = not settings.CI
900 if callable(progressbar):
901 progressbar = progressbar()
903 if isinstance(progressbar, bool) and progressbar:
904 progressbar = tqdm(
905 ncols=79,
906 ascii=bool(sys.platform == "win32"),
907 unit="B",
908 unit_scale=True,
909 leave=True,
910 )
912 if progressbar is not False:
913 progressbar.set_description(f"Downloading {extract_file_name(source)}")
915 headers: dict[str, str] = {}
916 if settings.user_agent is not None:
917 headers["User-Agent"] = settings.user_agent
918 elif settings.CI:
919 headers["User-Agent"] = "ci"
921 r = httpx.get(
922 str(source),
923 follow_redirects=True,
924 headers=headers,
925 timeout=settings.http_timeout,
926 )
927 _ = r.raise_for_status()
929 # set progressbar.total
930 total = r.headers.get("content-length")
931 if total is not None and not isinstance(total, int):
932 try:
933 total = int(total)
934 except Exception:
935 total = None
937 if progressbar is not False:
938 if total is None:
939 progressbar.total = 0
940 else:
941 progressbar.total = total
943 def iter_content():
944 for chunk in r.iter_bytes(chunk_size=4096):
945 yield chunk
946 if progressbar is not False:
947 _ = progressbar.update(len(chunk))
949 # Make sure the progress bar gets filled even if the actual number
950 # is chunks is smaller than expected. This happens when streaming
951 # text files that are compressed by the server when sending (gzip).
952 # Binary files don't experience this.
953 # (adapted from pooch.HttpDownloader)
954 if progressbar is not False:
955 progressbar.reset()
956 if total is not None:
957 _ = progressbar.update(total)
959 progressbar.close()
961 return iter_content()
964def extract_file_name(
965 src: ZarrSource | FileSource | FileDescr | ZipPath,
966) -> FileName:
967 if isinstance(src, FileDescr):
968 src = src.source
970 if isinstance(src, ZipPath):
971 return src.name or src.root.filename or "bioimageio.zip"
972 elif isinstance(src, (RelativeFilePath, RelativeDirectory)):
973 return src.path.name
974 elif isinstance(src, PurePath):
975 return src.name
976 else:
977 url = urlparse(str(src))
978 if (
979 url.scheme == "https"
980 and url.hostname == "zenodo.org"
981 and url.path.startswith("/api/records/")
982 and url.path.endswith("/content")
983 ):
984 return url.path.split("/")[-2]
985 else:
986 return url.path.split("/")[-1]
989def extract_file_descrs(
990 data: IncompleteDescrView,
991) -> list[FileDescr]:
992 collected: list[FileDescr] = []
993 with get_validation_context().replace(perform_io_checks=False, log_warnings=False):
994 _extract_file_descrs_impl(data, collected)
996 return collected
999def _extract_file_descrs_impl(
1000 data: IncompleteDescrView | IncompleteDescrInnerView,
1001 collected: list[FileDescr],
1002) -> None:
1003 if isinstance(data, FileDescr):
1004 collected.append(data)
1005 elif isinstance(data, Node):
1006 for _, v in data:
1007 _extract_file_descrs_impl(v, collected)
1008 elif isinstance(data, collections.abc.Mapping):
1009 if "source" in data and "sha256" in data:
1010 try:
1011 fd = FileDescr.model_validate(
1012 {"source": data["source"], "sha256": data["sha256"]}
1013 )
1014 except Exception:
1015 warnings.warn(
1016 "Found mapping with 'source' and 'sha256' keys, but could not parse it as a FileDescr. Ignoring `sha256`."
1017 )
1018 try:
1019 fd = FileDescr.model_validate({"source": data["source"]})
1020 except Exception:
1021 warnings.warn(
1022 f"Found mapping with 'source' and `sha256' keys , but could not parse it as a FileDescr, evning when ignoring 'sha256'. Ignoring `source`: {data['source']}."
1023 )
1024 else:
1025 collected.append(fd)
1026 else:
1027 collected.append(fd)
1029 for v in data.values():
1030 _extract_file_descrs_impl(v, collected)
1031 elif not isinstance(data, (str, Path, RelativeFilePath)) and isinstance(
1032 data, collections.abc.Sequence
1033 ):
1034 for v in data:
1035 _extract_file_descrs_impl(v, collected)
1038F = TypeVar("F", bound=Union[ZarrSource, FileSource, FileDescr])
1041def validate_suffix(
1042 value: F,
1043 suffix: str | Sequence[str],
1044 *,
1045 case_sensitive: bool,
1046 allow_any_parent_suffix: bool = False,
1047) -> F:
1048 """check final suffix"""
1049 if isinstance(suffix, str):
1050 suffixes = [suffix]
1051 else:
1052 suffixes = suffix
1054 assert len(suffixes) > 0, "no suffix given"
1055 assert all(suff.startswith(".") for suff in suffixes), (
1056 "expected suffixes to start with '.'"
1057 )
1058 o_value = value
1059 if isinstance(value, FileDescr):
1060 strict = value.source
1061 elif isinstance(value, (str, pydantic.AnyUrl)):
1062 strict = interprete_file_source(value)
1063 else:
1064 strict = value
1066 if isinstance(strict, (HttpUrl, pydantic.AnyUrl, RootHttpUrl, FtpUrl)):
1067 if strict.path is None or "." not in (path := strict.path):
1068 actual_suffixes = []
1069 else:
1070 if (
1071 strict.host == "zenodo.org"
1072 and path.startswith("/api/records/")
1073 and path.endswith("/content")
1074 ):
1075 # Zenodo API URLs have a "/content" suffix that should be ignored
1076 path = path[: -len("/content")]
1078 actual_suffixes = [f".{path.split('.')[-1]}"]
1080 elif isinstance(strict, PurePath):
1081 actual_suffixes = strict.suffixes
1082 elif isinstance(strict, (RelativeFilePath, RelativeDirectory)):
1083 actual_suffixes = strict.path.suffixes
1084 else:
1085 assert_never(strict)
1087 if actual_suffixes:
1088 actual_suffix = actual_suffixes[-1]
1089 else:
1090 actual_suffix = "no suffix"
1092 if (
1093 case_sensitive
1094 and actual_suffix not in suffixes
1095 or not case_sensitive
1096 and actual_suffix.lower() not in [s.lower() for s in suffixes]
1097 ):
1098 if allow_any_parent_suffix and strict != strict.parent:
1099 try:
1100 _ = validate_suffix(
1101 strict.parent,
1102 suffix,
1103 case_sensitive=case_sensitive,
1104 allow_any_parent_suffix=allow_any_parent_suffix,
1105 )
1106 except ValueError:
1107 pass
1108 else:
1109 return o_value
1111 if len(suffixes) == 1:
1112 raise ValueError(f"Expected suffix {suffixes[0]}, but got {actual_suffix}")
1113 else:
1114 raise ValueError(
1115 f"Expected a suffix from {suffixes}, but got {actual_suffix}"
1116 )
1118 return o_value
1121def populate_cache(sources: Sequence[FileDescr | LightHttpFileDescr]):
1122 unique: set[str] = set()
1123 for src in sources:
1124 if src.sha256 is None:
1125 continue # not caching without known SHA
1127 if isinstance(src.source, (HttpUrl, pydantic.AnyUrl)):
1128 url = str(src.source)
1129 elif isinstance(src.source, RelativeFilePath):
1130 if isinstance(absolute := src.source.absolute(), HttpUrl):
1131 url = str(absolute)
1132 else:
1133 continue # not caching local paths
1134 elif isinstance(src.source, Path):
1135 continue # not caching local paths
1136 elif isinstance(src.source, (RootHttpUrl, RelativeDirectory)):
1137 continue # not caching directories
1138 elif isinstance(src.source, FtpUrl):
1139 continue # not caching FTP URLs
1140 else:
1141 assert_never(src.source)
1143 if url in unique:
1144 continue # skip duplicate URLs
1146 unique.add(url)
1147 _ = src.download()